This is page 2 of 4. Use http://codebase.md/grafana/mcp-grafana?lines=false&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
--------------------------------------------------------------------------------
/tls_test.go:
--------------------------------------------------------------------------------
```go
package mcpgrafana
import (
"fmt"
"net/http"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestTLSConfig_CreateTLSConfig(t *testing.T) {
t.Run("nil config returns nil", func(t *testing.T) {
var config *TLSConfig
tlsCfg, err := config.CreateTLSConfig()
assert.NoError(t, err)
assert.Nil(t, tlsCfg)
})
t.Run("skip verify only", func(t *testing.T) {
config := &TLSConfig{SkipVerify: true}
tlsCfg, err := config.CreateTLSConfig()
assert.NoError(t, err)
require.NotNil(t, tlsCfg)
assert.True(t, tlsCfg.InsecureSkipVerify)
assert.Empty(t, tlsCfg.Certificates)
assert.Nil(t, tlsCfg.RootCAs)
})
t.Run("invalid cert file", func(t *testing.T) {
config := &TLSConfig{
CertFile: "nonexistent.pem",
KeyFile: "nonexistent.key",
}
_, err := config.CreateTLSConfig()
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to load client certificate")
})
t.Run("invalid CA file", func(t *testing.T) {
config := &TLSConfig{
CAFile: "nonexistent-ca.pem",
}
_, err := config.CreateTLSConfig()
assert.Error(t, err)
assert.Contains(t, err.Error(), "failed to read CA certificate")
})
}
func TestHTTPTransport(t *testing.T) {
t.Run("nil TLS config", func(t *testing.T) {
var tlsConfig *TLSConfig
transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
assert.NoError(t, err)
assert.NotNil(t, transport)
// Should be default transport clone
httpTransport, ok := transport.(*http.Transport)
require.True(t, ok)
assert.NotNil(t, httpTransport)
})
t.Run("skip verify config", func(t *testing.T) {
tlsConfig := &TLSConfig{SkipVerify: true}
transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
assert.NoError(t, err)
require.NotNil(t, transport)
httpTransport, ok := transport.(*http.Transport)
require.True(t, ok)
require.NotNil(t, httpTransport.TLSClientConfig)
assert.True(t, httpTransport.TLSClientConfig.InsecureSkipVerify)
})
t.Run("invalid TLS config", func(t *testing.T) {
tlsConfig := &TLSConfig{
CertFile: "nonexistent.pem",
KeyFile: "nonexistent.key",
}
_, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
assert.Error(t, err)
})
}
// mockRoundTripper is a mock implementation of http.RoundTripper for testing
type mockRoundTripper struct {
capturedRequest *http.Request
response *http.Response
err error
}
func (m *mockRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
m.capturedRequest = req
if m.response != nil {
return m.response, m.err
}
// Return a default successful response
return &http.Response{
StatusCode: 200,
Header: make(http.Header),
Body: http.NoBody,
}, m.err
}
func TestUserAgentTransport(t *testing.T) {
tests := []struct {
name string
userAgent string
existingUserAgent string
expectedUserAgent string
}{
{
name: "sets user agent when not present",
userAgent: "mcp-grafana/1.0.0",
existingUserAgent: "",
expectedUserAgent: "mcp-grafana/1.0.0",
},
{
name: "does not override existing user agent",
userAgent: "mcp-grafana/1.0.0",
existingUserAgent: "existing-client/2.0.0",
expectedUserAgent: "existing-client/2.0.0",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
// Create mock round tripper
mockRT := &mockRoundTripper{}
// Create user agent transport
transport := &UserAgentTransport{
rt: mockRT,
UserAgent: tt.userAgent,
}
// Create request
req, err := http.NewRequest("GET", "http://example.com", nil)
require.NoError(t, err)
// Set existing user agent if specified
if tt.existingUserAgent != "" {
req.Header.Set("User-Agent", tt.existingUserAgent)
}
// Make request through transport
_, err = transport.RoundTrip(req)
require.NoError(t, err)
// Verify user agent header
assert.Equal(t, tt.expectedUserAgent, mockRT.capturedRequest.Header.Get("User-Agent"))
})
}
}
func TestVersion(t *testing.T) {
version := Version()
assert.NotEmpty(t, version)
// Version should be either "(devel)" for development builds or a proper version
assert.True(t, version == "(devel)" || len(version) > 0)
}
func TestUserAgent(t *testing.T) {
userAgent := UserAgent()
assert.Contains(t, userAgent, "mcp-grafana/")
assert.NotEqual(t, "mcp-grafana/", userAgent) // Should have version appended
// Should match the pattern mcp-grafana/{version}
version := Version()
expected := fmt.Sprintf("mcp-grafana/%s", version)
assert.Equal(t, expected, userAgent)
}
func TestNewUserAgentTransport(t *testing.T) {
t.Run("with explicit user agent", func(t *testing.T) {
mockRT := &mockRoundTripper{}
userAgent := "test-agent/1.0.0"
transport := NewUserAgentTransport(mockRT, userAgent)
assert.Equal(t, mockRT, transport.rt)
assert.Equal(t, userAgent, transport.UserAgent)
})
t.Run("with default user agent", func(t *testing.T) {
mockRT := &mockRoundTripper{}
transport := NewUserAgentTransport(mockRT)
assert.Equal(t, mockRT, transport.rt)
assert.Equal(t, UserAgent(), transport.UserAgent)
assert.Contains(t, transport.UserAgent, "mcp-grafana/")
})
}
func TestNewUserAgentTransportWithNilRoundTripper(t *testing.T) {
t.Run("with explicit user agent", func(t *testing.T) {
userAgent := "test-agent/1.0.0"
transport := NewUserAgentTransport(nil, userAgent)
assert.Equal(t, http.DefaultTransport, transport.rt)
assert.Equal(t, userAgent, transport.UserAgent)
})
t.Run("with default user agent", func(t *testing.T) {
transport := NewUserAgentTransport(nil)
assert.Equal(t, http.DefaultTransport, transport.rt)
assert.Equal(t, UserAgent(), transport.UserAgent)
assert.Contains(t, transport.UserAgent, "mcp-grafana/")
})
}
```
--------------------------------------------------------------------------------
/tools/incident.go:
--------------------------------------------------------------------------------
```go
package tools
import (
"context"
"fmt"
"github.com/grafana/incident-go"
mcpgrafana "github.com/grafana/mcp-grafana"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
type ListIncidentsParams struct {
Limit int `json:"limit" jsonschema:"description=The maximum number of incidents to return"`
Drill bool `json:"drill" jsonschema:"description=Whether to include drill incidents"`
Status string `json:"status" jsonschema:"description=The status of the incidents to include. Valid values: 'active'\\, 'resolved'"`
}
func listIncidents(ctx context.Context, args ListIncidentsParams) (*incident.QueryIncidentPreviewsResponse, error) {
c := mcpgrafana.IncidentClientFromContext(ctx)
is := incident.NewIncidentsService(c)
// Set default limit to 10 if not specified
limit := args.Limit
if limit <= 0 {
limit = 10
}
query := ""
if !args.Drill {
query = "isdrill:false"
}
if args.Status != "" {
query += fmt.Sprintf(" status:%s", args.Status)
}
incidents, err := is.QueryIncidentPreviews(ctx, incident.QueryIncidentPreviewsRequest{
Query: incident.IncidentPreviewsQuery{
QueryString: query,
OrderDirection: "DESC",
Limit: limit,
},
})
if err != nil {
return nil, fmt.Errorf("list incidents: %w", err)
}
return incidents, nil
}
var ListIncidents = mcpgrafana.MustTool(
"list_incidents",
"List Grafana incidents. Allows filtering by status ('active', 'resolved') and optionally including drill incidents. Returns a preview list with basic details.",
listIncidents,
mcp.WithTitleAnnotation("List incidents"),
mcp.WithIdempotentHintAnnotation(true),
mcp.WithReadOnlyHintAnnotation(true),
)
type CreateIncidentParams struct {
Title string `json:"title" jsonschema:"description=The title of the incident"`
Severity string `json:"severity" jsonschema:"description=The severity of the incident"`
RoomPrefix string `json:"roomPrefix" jsonschema:"description=The prefix of the room to create the incident in"`
IsDrill bool `json:"isDrill" jsonschema:"description=Whether the incident is a drill incident"`
Status string `json:"status" jsonschema:"description=The status of the incident"`
AttachCaption string `json:"attachCaption" jsonschema:"description=The caption of the attachment"`
AttachURL string `json:"attachUrl" jsonschema:"description=The URL of the attachment"`
Labels []incident.IncidentLabel `json:"labels" jsonschema:"description=The labels to add to the incident"`
}
func createIncident(ctx context.Context, args CreateIncidentParams) (*incident.Incident, error) {
c := mcpgrafana.IncidentClientFromContext(ctx)
is := incident.NewIncidentsService(c)
incident, err := is.CreateIncident(ctx, incident.CreateIncidentRequest{
Title: args.Title,
Severity: args.Severity,
RoomPrefix: args.RoomPrefix,
IsDrill: args.IsDrill,
Status: args.Status,
AttachCaption: args.AttachCaption,
AttachURL: args.AttachURL,
Labels: args.Labels,
})
if err != nil {
return nil, fmt.Errorf("create incident: %w", err)
}
return &incident.Incident, nil
}
var CreateIncident = mcpgrafana.MustTool(
"create_incident",
"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.",
createIncident,
mcp.WithTitleAnnotation("Create incident"),
)
type AddActivityToIncidentParams struct {
IncidentID string `json:"incidentId" jsonschema:"description=The ID of the incident to add the activity to"`
Body string `json:"body" jsonschema:"description=The body of the activity. URLs will be parsed and attached as context"`
EventTime string `json:"eventTime" jsonschema:"description=The time that the activity occurred. If not provided\\, the current time will be used"`
}
func addActivityToIncident(ctx context.Context, args AddActivityToIncidentParams) (*incident.ActivityItem, error) {
c := mcpgrafana.IncidentClientFromContext(ctx)
as := incident.NewActivityService(c)
activity, err := as.AddActivity(ctx, incident.AddActivityRequest{
IncidentID: args.IncidentID,
ActivityKind: "userNote",
Body: args.Body,
EventTime: args.EventTime,
})
if err != nil {
return nil, fmt.Errorf("add activity to incident: %w", err)
}
return &activity.ActivityItem, nil
}
var AddActivityToIncident = mcpgrafana.MustTool(
"add_activity_to_incident",
"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.",
addActivityToIncident,
mcp.WithTitleAnnotation("Add activity to incident"),
)
func AddIncidentTools(mcp *server.MCPServer, enableWriteTools bool) {
ListIncidents.Register(mcp)
if enableWriteTools {
CreateIncident.Register(mcp)
AddActivityToIncident.Register(mcp)
}
GetIncident.Register(mcp)
}
type GetIncidentParams struct {
ID string `json:"id" jsonschema:"description=The ID of the incident to retrieve"`
}
func getIncident(ctx context.Context, args GetIncidentParams) (*incident.Incident, error) {
c := mcpgrafana.IncidentClientFromContext(ctx)
is := incident.NewIncidentsService(c)
incidentResp, err := is.GetIncident(ctx, incident.GetIncidentRequest{
IncidentID: args.ID,
})
if err != nil {
return nil, fmt.Errorf("get incident by ID: %w", err)
}
return &incidentResp.Incident, nil
}
var GetIncident = mcpgrafana.MustTool(
"get_incident",
"Get a single incident by ID. Returns the full incident details including title, status, severity, labels, timestamps, and other metadata.",
getIncident,
mcp.WithTitleAnnotation("Get incident details"),
mcp.WithIdempotentHintAnnotation(true),
mcp.WithReadOnlyHintAnnotation(true),
)
```
--------------------------------------------------------------------------------
/examples/tls_example.go:
--------------------------------------------------------------------------------
```go
package main
import (
"context"
"fmt"
"log"
"net/http"
"os"
mcpgrafana "github.com/grafana/mcp-grafana"
"github.com/grafana/mcp-grafana/tools"
"github.com/mark3labs/mcp-go/server"
)
func main() {
// Example 1: Basic TLS configuration with skip verify (for testing)
fmt.Println("Example 1: Basic TLS configuration with skip verify")
basicTLSExample()
// Example 2: Full mTLS configuration with client certificates
fmt.Println("\nExample 2: Full mTLS configuration")
fullTLSExample()
// Example 3: Running an MCP server with TLS support
fmt.Println("\nExample 3: MCP server with TLS support")
if len(os.Args) > 1 && os.Args[1] == "run-server" {
runServerWithTLS()
} else {
fmt.Println("Use 'go run tls_example.go run-server' to actually start the server")
showServerExample()
}
}
func basicTLSExample() {
// Create a TLS config that skips certificate verification
// This is useful for testing against self-signed certificates
tlsConfig := &mcpgrafana.TLSConfig{SkipVerify: true}
// Create a Grafana config with TLS support
grafanaConfig := mcpgrafana.GrafanaConfig{
Debug: true,
TLSConfig: tlsConfig,
}
// Create a context function that includes TLS configuration
contextFunc := mcpgrafana.ComposedStdioContextFunc(grafanaConfig)
// Test the context function
ctx := contextFunc(context.Background())
// Verify the configuration is applied
retrievedConfig := mcpgrafana.GrafanaConfigFromContext(ctx)
if retrievedConfig.TLSConfig != nil {
fmt.Printf("✓ TLS configuration applied: SkipVerify=%v\n", retrievedConfig.TLSConfig.SkipVerify)
}
fmt.Printf("✓ Debug mode enabled: %v\n", retrievedConfig.Debug)
}
func fullTLSExample() {
// Example paths for certificate files
// In a real scenario, these would point to actual certificate files
certFile := "/path/to/client.crt"
keyFile := "/path/to/client.key"
caFile := "/path/to/ca.crt"
// Create TLS config with client certificates and CA verification
tlsConfig := &mcpgrafana.TLSConfig{
CertFile: certFile,
KeyFile: keyFile,
CAFile: caFile,
}
// Create Grafana config with TLS support
grafanaConfig := mcpgrafana.GrafanaConfig{
Debug: false,
TLSConfig: tlsConfig,
}
fmt.Printf("✓ TLS configuration created:\n")
fmt.Printf(" - Client cert: %s\n", tlsConfig.CertFile)
fmt.Printf(" - Client key: %s\n", tlsConfig.KeyFile)
fmt.Printf(" - CA file: %s\n", tlsConfig.CAFile)
fmt.Printf(" - Skip verify: %v\n", tlsConfig.SkipVerify)
fmt.Printf(" - Debug mode: %v\n", grafanaConfig.Debug)
// Create context functions for different transport types
stdioFunc := mcpgrafana.ComposedStdioContextFunc(grafanaConfig)
sseFunc := mcpgrafana.ComposedSSEContextFunc(grafanaConfig)
httpFunc := mcpgrafana.ComposedHTTPContextFunc(grafanaConfig)
fmt.Printf("✓ Context functions created for all transport types\n")
_ = stdioFunc
_ = sseFunc
_ = httpFunc
}
func showServerExample() {
fmt.Println("Example MCP server configuration with TLS:")
fmt.Println(`// Create TLS configuration
tlsConfig := &mcpgrafana.TLSConfig{
CertFile: "/path/to/client.crt",
KeyFile: "/path/to/client.key",
CAFile: "/path/to/ca.crt",
}
// Create Grafana configuration
grafanaConfig := mcpgrafana.GrafanaConfig{
Debug: true,
TLSConfig: tlsConfig,
}
// Create MCP server
s := server.NewMCPServer("mcp-grafana", "1.0.0")
// Add tools
tools.AddSearchTools(s)
tools.AddDatasourceTools(s)
// ... add other tools as needed
// Create stdio server with TLS support
srv := server.NewStdioServer(s)
srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(grafanaConfig))
// Start server
srv.Listen(ctx, os.Stdin, os.Stdout)`)
}
func runServerWithTLS() {
// Set up environment variables (in practice, these would be set externally)
if os.Getenv("GRAFANA_URL") == "" {
if err := os.Setenv("GRAFANA_URL", "https://localhost:3000"); err != nil {
log.Printf("Failed to set GRAFANA_URL: %v", err)
}
}
// Check for service account token first, then fall back to deprecated API key
if os.Getenv("GRAFANA_SERVICE_ACCOUNT_TOKEN") == "" {
if os.Getenv("GRAFANA_API_KEY") == "" {
fmt.Println("Warning: Neither GRAFANA_SERVICE_ACCOUNT_TOKEN nor GRAFANA_API_KEY is set")
} else {
fmt.Println("Warning: GRAFANA_API_KEY is deprecated, please use GRAFANA_SERVICE_ACCOUNT_TOKEN instead")
}
}
// Create TLS configuration that skips verification for demo purposes
// In production, you would use real certificates
tlsConfig := &mcpgrafana.TLSConfig{SkipVerify: true}
grafanaConfig := mcpgrafana.GrafanaConfig{
Debug: true,
TLSConfig: tlsConfig,
}
// Create MCP server
s := server.NewMCPServer("mcp-grafana-tls-example", "1.0.0")
// Add some basic tools
tools.AddSearchTools(s)
tools.AddDatasourceTools(s)
tools.AddDashboardTools(s, false) // Read-only mode (no write tools)
// Create stdio server with TLS-enabled context function
srv := server.NewStdioServer(s)
srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(grafanaConfig))
fmt.Printf("Starting MCP Grafana server with TLS support...\n")
fmt.Printf("Grafana URL: %s\n", os.Getenv("GRAFANA_URL"))
fmt.Printf("TLS Skip Verify: %v\n", tlsConfig.SkipVerify)
// Start the server
ctx := context.Background()
if err := srv.Listen(ctx, os.Stdin, os.Stdout); err != nil {
log.Fatalf("Server error: %v", err)
}
}
// Example of creating custom HTTP clients with TLS configuration
func customClientExample() { //nolint:unused // Example function for documentation
ctx := context.Background()
// Add Grafana configuration to context
tlsConfig := &mcpgrafana.TLSConfig{
CertFile: "/path/to/cert.pem",
KeyFile: "/path/to/key.pem",
CAFile: "/path/to/ca.pem",
}
config := mcpgrafana.GrafanaConfig{
TLSConfig: tlsConfig,
}
ctx = mcpgrafana.WithGrafanaConfig(ctx, config)
_ = ctx // Use ctx to avoid ineffectual assignment warning
// Create custom HTTP transport with TLS
transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
if err != nil {
log.Fatalf("Failed to create transport: %v", err)
}
// Use the transport in your HTTP client
_ = transport
fmt.Println("✓ Custom HTTP transport created with TLS configuration")
}
```
--------------------------------------------------------------------------------
/internal/linter/jsonschema/jsonschema_lint.go:
--------------------------------------------------------------------------------
```go
package linter
import (
"fmt"
"go/ast"
"go/parser"
"go/token"
"os"
"path/filepath"
"regexp"
"sort"
"strings"
)
// JSONSchemaLinter checks for unescaped commas in jsonschema struct tags
type JSONSchemaLinter struct {
FilePaths []string
Errors []JSONSchemaError
FixMode bool
Fixed map[string]bool
}
// JSONSchemaError represents a linting error with file position details
type JSONSchemaError struct {
FilePath string
Line int
Column int
Offset int // Byte offset in the file
Struct string
Field string
Tag string
FixedTag string
}
// tagPattern matches jsonschema tags with description containing unescaped commas
// It captures:
// 1. The jsonschema tag
// 2. Parts of the description containing unescaped commas
// The pattern correctly handles:
// - Simple unescaped comma: "description=Something, with comma"
// - Escaped quote followed by unescaped comma: "description=With \"quote, and comma"
// - But not match escaped comma: "description=With escaped\, comma"
var tagPattern = regexp.MustCompile(`jsonschema:"([^"]*)description=(.*?[^\\],)([^"]*)"`)
// FindUnescapedCommas scans Go files for jsonschema struct tags with unescaped commas in descriptions
func (l *JSONSchemaLinter) FindUnescapedCommas(baseDir string) error {
// Reset errors
l.Errors = nil
if l.FixMode {
l.Fixed = make(map[string]bool)
}
// Walk through the directory and find Go files
err := filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error {
if err != nil {
return err
}
// Skip non-Go files
if !info.IsDir() && strings.HasSuffix(path, ".go") {
l.FilePaths = append(l.FilePaths, path)
}
return nil
})
if err != nil {
return fmt.Errorf("error walking directory: %v", err)
}
// Parse all Go files and check for the unescaped commas
for _, path := range l.FilePaths {
fset := token.NewFileSet()
f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
if err != nil {
return fmt.Errorf("error parsing file %s: %v", path, err)
}
fileErrors := []JSONSchemaError{}
// Visit all struct types
ast.Inspect(f, func(n ast.Node) bool {
ts, ok := n.(*ast.TypeSpec)
if !ok || ts.Type == nil {
return true
}
st, ok := ts.Type.(*ast.StructType)
if !ok {
return true
}
structName := ts.Name.Name
// Check each field of the struct
for _, field := range st.Fields.List {
if field.Tag == nil {
continue
}
tag := field.Tag.Value
// Check if the tag has a jsonschema description with unescaped comma
matches := tagPattern.FindStringSubmatch(tag)
if len(matches) > 0 {
fieldName := ""
if len(field.Names) > 0 {
fieldName = field.Names[0].Name
}
// Generate the fixed tag by escaping the commas in the description
fixedTag := tag
if len(matches) > 2 {
descWithUnescapedCommas := matches[2]
// Escape all unescaped commas
fixedDesc := escapeUnescapedCommas(descWithUnescapedCommas)
// Replace the original description with the fixed one
fixedTag = strings.Replace(tag, descWithUnescapedCommas, fixedDesc, 1)
}
pos := fset.Position(field.Tag.Pos())
errorInfo := JSONSchemaError{
FilePath: path,
Line: pos.Line,
Column: pos.Column,
Offset: pos.Offset,
Struct: structName,
Field: fieldName,
Tag: tag,
FixedTag: fixedTag,
}
fileErrors = append(fileErrors, errorInfo)
}
}
return true
})
// Add all errors for this file
l.Errors = append(l.Errors, fileErrors...)
// If in fix mode and we found errors, fix the file
if l.FixMode && len(fileErrors) > 0 {
err := l.fixFile(path, fileErrors)
if err != nil {
return fmt.Errorf("error fixing file %s: %v", path, err)
}
l.Fixed[path] = true
}
}
return nil
}
// escapeUnescapedCommas escapes any unescaped commas in the description
func escapeUnescapedCommas(desc string) string {
// Use regex to find all commas that are not preceded by a backslash
r := regexp.MustCompile(`([^\\]),`)
// Replace them with the same text but with an escaped comma
return r.ReplaceAllString(desc, `$1\\,`)
}
// fixFile applies the fixes to a file
func (l *JSONSchemaLinter) fixFile(path string, errors []JSONSchemaError) error {
// Read the file content
content, err := os.ReadFile(path)
if err != nil {
return fmt.Errorf("error reading file %s: %v", path, err)
}
// Convert to string for easier manipulation
fileContent := string(content)
// Sort errors by offset in reverse order to avoid offset changes
sort.Slice(errors, func(i, j int) bool {
return errors[i].Offset > errors[j].Offset
})
// Apply fixes
for _, e := range errors {
// Find the tag in the file content
tagStart := strings.Index(fileContent[e.Offset:], e.Tag)
if tagStart == -1 {
continue
}
absOffset := e.Offset + tagStart
// Replace the tag with the fixed version
fixedContent := fileContent[:absOffset] + e.FixedTag + fileContent[absOffset+len(e.Tag):]
fileContent = fixedContent
}
// Write back to the file
err = os.WriteFile(path, []byte(fileContent), 0644)
if err != nil {
return fmt.Errorf("error writing file %s: %v", path, err)
}
return nil
}
// PrintErrors outputs all the found errors
func (l *JSONSchemaLinter) PrintErrors() {
if len(l.Errors) == 0 {
fmt.Println("No unescaped commas found in jsonschema descriptions.")
return
}
if l.FixMode {
fmt.Printf("Found and fixed %d unescaped commas in jsonschema descriptions:\n\n", len(l.Errors))
} else {
fmt.Printf("Found %d unescaped commas in jsonschema descriptions:\n\n", len(l.Errors))
}
for i, err := range l.Errors {
relPath, _ := filepath.Rel(".", err.FilePath)
fmt.Printf("%d. %s:%d:%d - Struct: %s, Field: %s\n",
i+1, relPath, err.Line, err.Column, err.Struct, err.Field)
fmt.Printf(" - %s\n", err.Tag)
if l.FixMode {
fmt.Printf(" - Fixed to: %s\n\n", err.FixedTag)
} else {
fmt.Printf(" - Commas in description must be escaped with \\\\,\n\n")
}
}
if !l.FixMode {
fmt.Println("Please escape all commas in jsonschema descriptions with \\\\, to prevent truncation.")
fmt.Println("You can run with --fix to automatically fix these issues.")
} else {
fixedFileCount := len(l.Fixed)
fmt.Printf("Fixed %d file(s).\n", fixedFileCount)
}
}
```
--------------------------------------------------------------------------------
/tests/navigation_test.py:
--------------------------------------------------------------------------------
```python
import json
import pytest
from langevals import expect
from langevals_langevals.llm_boolean import (
CustomLLMBooleanEvaluator,
CustomLLMBooleanSettings,
)
from litellm import Message, acompletion
from mcp import ClientSession
from mcp.types import TextContent
from conftest import models
from utils import (
get_converted_tools,
llm_tool_call_sequence,
flexible_tool_call,
)
pytestmark = pytest.mark.anyio
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_generate_dashboard_deeplink(model: str, mcp_client: ClientSession):
tools = await get_converted_tools(mcp_client)
prompt = """Please create a dashboard deeplink for dashboard with UID 'test-uid'."""
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
messages = await llm_tool_call_sequence(
model, messages, tools, mcp_client, "generate_deeplink",
{"resourceType": "dashboard", "dashboardUid": "test-uid"}
)
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
assert "/d/test-uid" in content, f"Expected dashboard URL with /d/test-uid, got: {content}"
dashboard_link_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response contain a URL with /d/ path and the dashboard UID?",
)
)
print("Dashboard deeplink content:", content)
expect(input=prompt, output=content).to_pass(dashboard_link_checker)
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_generate_panel_deeplink(model: str, mcp_client: ClientSession):
tools = await get_converted_tools(mcp_client)
prompt = "Generate a deeplink for panel 5 in dashboard with UID 'test-uid'"
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
messages = await llm_tool_call_sequence(
model, messages, tools, mcp_client, "generate_deeplink",
{
"resourceType": "panel",
"dashboardUid": "test-uid",
"panelId": 5
}
)
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
assert "viewPanel=5" in content, f"Expected panel URL with viewPanel=5, got: {content}"
panel_link_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response contain a URL with viewPanel parameter?",
)
)
print("Panel deeplink content:", content)
expect(input=prompt, output=content).to_pass(panel_link_checker)
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_generate_explore_deeplink(model: str, mcp_client: ClientSession):
tools = await get_converted_tools(mcp_client)
prompt = "Generate a deeplink for Grafana Explore with datasource 'test-uid'"
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
messages = await llm_tool_call_sequence(
model, messages, tools, mcp_client, "generate_deeplink",
{"resourceType": "explore", "datasourceUid": "test-uid"}
)
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
assert "/explore" in content, f"Expected explore URL with /explore path, got: {content}"
explore_link_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response contain a URL with /explore path?",
)
)
print("Explore deeplink content:", content)
expect(input=prompt, output=content).to_pass(explore_link_checker)
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_generate_deeplink_with_time_range(model: str, mcp_client: ClientSession):
tools = await get_converted_tools(mcp_client)
prompt = "Generate a dashboard deeplink for 'test-uid' showing the last 6 hours"
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
messages = await llm_tool_call_sequence(
model, messages, tools, mcp_client, "generate_deeplink",
{
"resourceType": "dashboard",
"dashboardUid": "test-uid",
"timeRange": {
"from": "now-6h",
"to": "now"
}
}
)
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
assert "from=now-6h" in content and "to=now" in content, f"Expected time range parameters, got: {content}"
time_range_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response contain a URL with time range parameters?",
)
)
print("Time range deeplink content:", content)
expect(input=prompt, output=content).to_pass(time_range_checker)
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_generate_deeplink_with_query_params(model: str, mcp_client: ClientSession):
tools = await get_converted_tools(mcp_client)
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"
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
# Use flexible tool call with required parameters
messages = await flexible_tool_call(
model, messages, tools, mcp_client, "generate_deeplink",
required_params={"resourceType": "dashboard", "dashboardUid": "test-uid"}
)
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
# Verify both specific query parameters are in the final URL
assert "var-datasource=prometheus" in content, f"Expected var-datasource=prometheus in URL, got: {content}"
assert "refresh=30s" in content, f"Expected refresh=30s in URL, got: {content}"
custom_params_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response contain a URL with custom query parameters?",
)
)
print("Custom params deeplink content:", content)
expect(input=prompt, output=content).to_pass(custom_params_checker)
```
--------------------------------------------------------------------------------
/tests/utils.py:
--------------------------------------------------------------------------------
```python
import json
from litellm.types.utils import ModelResponse
from litellm import acompletion, Choices, Message
from mcp.types import TextContent, Tool
async def assert_and_handle_tool_call(
response: ModelResponse,
mcp_client,
expected_tool: str,
expected_args: dict = None,
) -> list:
messages = []
tool_calls = []
for c in response.choices:
assert isinstance(c, Choices)
tool_calls.extend(c.message.tool_calls or [])
messages.append(c.message)
# Better error message if wrong number of tool calls
if len(tool_calls) != 1:
actual_calls = [tc.function.name for tc in tool_calls] if tool_calls else []
assert len(tool_calls) == 1, (
f"\n❌ Expected exactly 1 tool call, got {len(tool_calls)}\n"
f"Expected tool: {expected_tool}\n"
f"Actual tools called: {actual_calls}\n"
f"LLM response: {response.choices[0].message.content if response.choices else 'N/A'}"
)
for tool_call in tool_calls:
actual_tool = tool_call.function.name
if actual_tool != expected_tool:
# Parse arguments to understand what LLM was trying to do
try:
actual_args = (
json.loads(tool_call.function.arguments)
if tool_call.function.arguments
else {}
)
except:
actual_args = tool_call.function.arguments
assert False, (
f"\n❌ LLM called wrong tool!\n"
f"Expected: {expected_tool}\n"
f"Got: {actual_tool}\n"
f"With args: {json.dumps(actual_args, indent=2)}\n"
f"\n💡 Debugging tips:\n"
f" - Check if the prompt clearly indicates which tool to use\n"
f" - Verify the expected tool exists in the available tools\n"
f" - Consider if the tool description is clear enough\n"
)
arguments = (
{}
if len(tool_call.function.arguments) == 0
else json.loads(tool_call.function.arguments)
)
if expected_args:
for key, value in expected_args.items():
if key not in arguments:
assert False, (
f"\n❌ Missing expected parameter '{key}'\n"
f"Expected args: {json.dumps(expected_args, indent=2)}\n"
f"Actual args: {json.dumps(arguments, indent=2)}\n"
)
if arguments[key] != value:
assert False, (
f"\n❌ Wrong value for parameter '{key}'\n"
f"Expected: {value}\n"
f"Got: {arguments[key]}\n"
f"Full args: {json.dumps(arguments, indent=2)}\n"
)
result = await mcp_client.call_tool(tool_call.function.name, arguments)
assert len(result.content) == 1, (
f"Expected one result for tool {tool_call.function.name}, got {len(result.content)}"
)
assert isinstance(result.content[0], TextContent), (
f"Expected TextContent for tool {tool_call.function.name}, got {type(result.content[0])}"
)
messages.append(
Message(
role="tool", tool_call_id=tool_call.id, content=result.content[0].text
)
)
return messages
def convert_tool(tool: Tool) -> dict:
return {
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": {
**tool.inputSchema,
"properties": tool.inputSchema.get("properties", {}),
},
},
}
async def llm_tool_call_sequence(
model, messages, tools, mcp_client, tool_name, tool_args=None
):
print(f"\n🤖 Calling LLM ({model}) and expecting tool: {tool_name}")
print(f"📝 Last message: {messages[-1].get('content', messages[-1])[:200]}...")
response = await acompletion(
model=model,
messages=messages,
tools=tools,
)
assert isinstance(response, ModelResponse)
# Print what tool was actually called for debugging
if response.choices and response.choices[0].message.tool_calls:
actual_tool = response.choices[0].message.tool_calls[0].function.name
print(f"✅ LLM called: {actual_tool}")
if actual_tool != tool_name:
print(f"⚠️ WARNING: Expected {tool_name} but got {actual_tool}")
messages.extend(
await assert_and_handle_tool_call(
response, mcp_client, tool_name, tool_args or {}
)
)
return messages
async def get_converted_tools(mcp_client):
tools = await mcp_client.list_tools()
return [convert_tool(t) for t in tools.tools]
async def flexible_tool_call(
model, messages, tools, mcp_client, expected_tool_name, required_params=None
):
"""
Make a flexible tool call that only checks essential parameters.
Returns updated messages list.
Args:
model: The LLM model to use
messages: Current conversation messages
tools: Available tools
mcp_client: MCP client session
expected_tool_name: Name of the tool we expect to be called
required_params: Dict of essential parameters to check (optional)
Returns:
Updated messages list including tool call and result
"""
response = await acompletion(model=model, messages=messages, tools=tools)
# Check that a tool call was made
assert response.choices[0].message.tool_calls is not None, (
f"Expected tool call for {expected_tool_name}"
)
assert len(response.choices[0].message.tool_calls) >= 1, (
f"Expected at least one tool call for {expected_tool_name}"
)
tool_call = response.choices[0].message.tool_calls[0]
assert tool_call.function.name == expected_tool_name, (
f"Expected {expected_tool_name} tool, got {tool_call.function.name}"
)
arguments = json.loads(tool_call.function.arguments)
# Check required parameters if specified
if required_params:
for key, expected_value in required_params.items():
assert key in arguments, f"Expected parameter '{key}' in tool arguments"
if expected_value is not None:
assert arguments[key] == expected_value, (
f"Expected {key}='{expected_value}', got {key}='{arguments.get(key)}'"
)
# Call the tool to verify it works
result = await mcp_client.call_tool(tool_call.function.name, arguments)
assert len(result.content) == 1
assert isinstance(result.content[0], TextContent)
# Add both the tool call and result to message history
messages.append(response.choices[0].message)
messages.append(
Message(role="tool", tool_call_id=tool_call.id, content=result.content[0].text)
)
return messages
```
--------------------------------------------------------------------------------
/tools/annotations_unit_test.go:
--------------------------------------------------------------------------------
```go
//go:build unit
package tools
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"net/url"
"strconv"
"testing"
"github.com/grafana/grafana-openapi-client-go/client"
"github.com/grafana/grafana-openapi-client-go/models"
mcpgrafana "github.com/grafana/mcp-grafana"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func mockCtxWithClient(server *httptest.Server) context.Context {
u, _ := url.Parse(server.URL)
cfg := client.DefaultTransportConfig()
cfg.Host = u.Host
cfg.Schemes = []string{"http"}
cfg.APIKey = "test"
c := client.NewHTTPClientWithConfig(nil, cfg)
return mcpgrafana.WithGrafanaClient(context.Background(), c)
}
func TestGetAnnotations_UsesCorrectQueryParams(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations", r.URL.Path)
q := r.URL.Query()
assert.Equal(t, "50", q.Get("limit"))
assert.Equal(t, "dash-1", q.Get("dashboardUID"))
assert.Equal(t, "true", q.Get("matchAny"))
assert.Equal(t, "tagA", q["tags"][0])
assert.Equal(t, "tagB", q["tags"][1])
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_ = json.NewEncoder(w).Encode([]interface{}{})
}))
defer server.Close()
ctx := mockCtxWithClient(server)
limit := int64(50)
uid := "dash-1"
matchAny := true
_, err := getAnnotations(ctx, GetAnnotationsInput{
Limit: &limit,
DashboardUID: &uid,
MatchAny: &matchAny,
Tags: []string{"tagA", "tagB"},
})
require.NoError(t, err)
}
func TestGetAnnotations_PropagatesError(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(`oops`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := getAnnotations(ctx, GetAnnotationsInput{})
require.Error(t, err)
assert.Contains(t, err.Error(), "get annotations:")
}
func TestCreateAnnotationGraphiteFormat_SendsCorrectBody_Minimal(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations/graphite", r.URL.Path)
assert.Equal(t, "POST", r.Method)
assert.Equal(t, "Bearer test", r.Header.Get("Authorization"))
var body models.PostGraphiteAnnotationsCmd
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
assert.Equal(t, "deploy", body.What)
assert.Equal(t, int64(1710000000000), body.When)
assert.Nil(t, body.Tags)
assert.Empty(t, body.Data)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"message":"annotation created"}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
What: "deploy",
When: 1710000000000,
})
require.NoError(t, err)
}
func TestCreateAnnotationGraphiteFormat_SendsCorrectBody_WithTagsAndData(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations/graphite", r.URL.Path)
assert.Equal(t, "POST", r.Method)
var body map[string]interface{}
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
assert.Equal(t, "incident", body["what"])
assert.Equal(t, float64(1720000000000), body["when"])
assert.ElementsMatch(t, []interface{}{"sev1", "network"}, body["tags"].([]interface{}))
assert.Equal(t, "context", body["data"])
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"message":"ok"}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
What: "incident",
When: 1720000000000,
Tags: []string{"sev1", "network"},
Data: "context",
})
require.NoError(t, err)
}
func TestCreateAnnotation_SendsCorrectBody(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations", r.URL.Path)
assert.Equal(t, "POST", r.Method)
var body models.PostAnnotationsCmd
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
assert.Equal(t, int64(7), body.PanelID)
assert.Equal(t, "hello", *body.Text)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"id": 1}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := createAnnotation(ctx, CreateAnnotationInput{
PanelID: 7,
Text: "hello",
})
require.NoError(t, err)
}
func TestCreateAnnotation_ErrorWrapped(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := createAnnotation(ctx, CreateAnnotationInput{Text: "t"})
require.Error(t, err)
assert.Contains(t, err.Error(), "create annotation:")
}
func TestCreateAnnotationGraphiteFormat_HTTPError(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusInternalServerError)
_, _ = w.Write([]byte(`internal error`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
What: "bad",
When: 1700000000000,
})
require.Error(t, err)
assert.Contains(t, err.Error(), "create graphite annotation")
}
func TestUpdateAnnotation_SendsCorrectBodyAndPath(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations/"+strconv.Itoa(55), r.URL.Path)
assert.Equal(t, "PUT", r.Method)
var body models.UpdateAnnotationsCmd
_ = json.NewDecoder(r.Body).Decode(&body)
assert.Equal(t, int64(111), body.Time)
assert.Equal(t, int64(222), body.TimeEnd)
assert.Equal(t, "hello", body.Text)
assert.Equal(t, []string{"a", "b"}, body.Tags)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
_, err := updateAnnotation(ctx, UpdateAnnotationInput{
ID: 55,
Time: 111,
TimeEnd: 222,
Text: "hello",
Tags: []string{"a", "b"},
})
require.NoError(t, err)
}
func TestPatchAnnotation_SendsOnlyProvidedFields(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations/"+strconv.Itoa(9), r.URL.Path)
assert.Equal(t, "PATCH", r.Method)
var body map[string]interface{}
_ = json.NewDecoder(r.Body).Decode(&body)
assert.Equal(t, "patched", body["text"])
assert.ElementsMatch(t, []interface{}{"x"}, body["tags"].([]interface{}))
assert.Nil(t, body["time"])
assert.Nil(t, body["timeEnd"])
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
text := "patched"
_, err := patchAnnotation(ctx, PatchAnnotationInput{
ID: 9,
Text: &text,
Tags: []string{"x"},
})
require.NoError(t, err)
}
func TestGetAnnotationTags_UsesCorrectQueryParams(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
assert.Equal(t, "/api/annotations/tags", r.URL.Path)
q := r.URL.Query()
assert.Equal(t, "error", q.Get("tag"))
assert.Equal(t, "50", q.Get("limit"))
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`{"result":{"tags":[]}}`))
}))
defer server.Close()
ctx := mockCtxWithClient(server)
tag := "error"
limit := "50"
_, err := getAnnotationTags(ctx, GetAnnotationTagsInput{
Tag: &tag,
Limit: &limit,
})
require.NoError(t, err)
}
```
--------------------------------------------------------------------------------
/tools/alerting_unit_test.go:
--------------------------------------------------------------------------------
```go
package tools
import (
"testing"
"github.com/stretchr/testify/require"
)
// Unit tests for parameter validation (no integration tag needed)
func TestCreateAlertRuleParams_Validate(t *testing.T) {
t.Run("valid parameters", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.NoError(t, err)
})
t.Run("missing title", func(t *testing.T) {
params := CreateAlertRuleParams{
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "title is required")
})
t.Run("missing rule group", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "ruleGroup is required")
})
t.Run("missing folder UID", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "folderUID is required")
})
t.Run("missing condition", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "condition is required")
})
t.Run("missing data", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "data is required")
})
t.Run("missing no data state", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "noDataState is required")
})
t.Run("missing exec error state", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "execErrState is required")
})
t.Run("missing for duration", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "for duration is required")
})
t.Run("invalid org ID", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 0,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "orgID is required and must be greater than 0")
})
}
func TestUpdateAlertRuleParams_Validate(t *testing.T) {
t.Run("valid parameters", func(t *testing.T) {
params := UpdateAlertRuleParams{
UID: "test-uid",
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.NoError(t, err)
})
t.Run("missing UID", func(t *testing.T) {
params := UpdateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "uid is required")
})
t.Run("invalid org ID", func(t *testing.T) {
params := UpdateAlertRuleParams{
UID: "test-uid",
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: -1,
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "orgID is required and must be greater than 0")
})
}
func TestDeleteAlertRuleParams_Validate(t *testing.T) {
t.Run("valid parameters", func(t *testing.T) {
params := DeleteAlertRuleParams{
UID: "test-uid",
}
err := params.validate()
require.NoError(t, err)
})
t.Run("missing UID", func(t *testing.T) {
params := DeleteAlertRuleParams{
UID: "",
}
err := params.validate()
require.Error(t, err)
require.Contains(t, err.Error(), "uid is required")
})
}
func TestBuiltInValidationCatchesInvalidData(t *testing.T) {
t.Run("invalid NoDataState enum value", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "InvalidValue", // Invalid enum
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
// Our simple validation won't catch this, but it would fail at API call
err := params.validate()
require.NoError(t, err, "Simple validation doesn't check enum values")
})
t.Run("invalid ExecErrState enum value", func(t *testing.T) {
params := CreateAlertRuleParams{
Title: "Test Rule",
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "BadValue", // Invalid enum
For: "5m",
OrgID: 1,
}
// Our simple validation won't catch this
err := params.validate()
require.NoError(t, err, "Simple validation doesn't check enum values")
})
t.Run("title too long", func(t *testing.T) {
longTitle := make([]byte, 200) // Max is 190
for i := range longTitle {
longTitle[i] = 'A'
}
params := CreateAlertRuleParams{
Title: string(longTitle),
RuleGroup: "test-group",
FolderUID: "test-folder",
Condition: "A",
Data: []interface{}{map[string]interface{}{"refId": "A"}},
NoDataState: "OK",
ExecErrState: "OK",
For: "5m",
OrgID: 1,
}
// Simple validation only checks if title is empty, not length
err := params.validate()
require.NoError(t, err, "Simple validation doesn't check length constraints")
})
}
```
--------------------------------------------------------------------------------
/tools.go:
--------------------------------------------------------------------------------
```go
package mcpgrafana
import (
"context"
"encoding/json"
"errors"
"fmt"
"reflect"
"github.com/invopop/jsonschema"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
)
// Tool represents a tool definition and its handler function for the MCP server.
// It encapsulates both the tool metadata (name, description, schema) and the function that executes when the tool is called.
// The simplest way to create a Tool is to use MustTool for compile-time tool creation,
// or ConvertTool if you need runtime tool creation with proper error handling.
type Tool struct {
Tool mcp.Tool
Handler server.ToolHandlerFunc
}
// Register adds the Tool to the given MCPServer.
// It is a convenience method that calls server.MCPServer.AddTool with the Tool's metadata and handler,
// allowing fluent tool registration in a single statement:
//
// mcpgrafana.MustTool(name, description, toolHandler).Register(server)
func (t *Tool) Register(mcp *server.MCPServer) {
mcp.AddTool(t.Tool, t.Handler)
}
// MustTool creates a new Tool from the given name, description, and toolHandler.
// It panics if the tool cannot be created, making it suitable for compile-time tool definitions where creation errors indicate programming mistakes.
func MustTool[T any, R any](
name, description string,
toolHandler ToolHandlerFunc[T, R],
options ...mcp.ToolOption,
) Tool {
tool, handler, err := ConvertTool(name, description, toolHandler, options...)
if err != nil {
panic(err)
}
return Tool{Tool: tool, Handler: handler}
}
// ToolHandlerFunc is the type of a handler function for a tool.
// 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.
type ToolHandlerFunc[T any, R any] = func(ctx context.Context, request T) (R, error)
// ConvertTool converts a toolHandler function to an MCP Tool and ToolHandlerFunc.
// The toolHandler must accept a context.Context and a struct with jsonschema tags for parameter documentation.
// The struct fields define the tool's input schema, while the return value can be a string, struct, or *mcp.CallToolResult.
// This function automatically generates JSON schema from the struct type and wraps the handler with OpenTelemetry instrumentation.
func ConvertTool[T any, R any](name, description string, toolHandler ToolHandlerFunc[T, R], options ...mcp.ToolOption) (mcp.Tool, server.ToolHandlerFunc, error) {
zero := mcp.Tool{}
handlerValue := reflect.ValueOf(toolHandler)
handlerType := handlerValue.Type()
if handlerType.Kind() != reflect.Func {
return zero, nil, errors.New("tool handler must be a function")
}
if handlerType.NumIn() != 2 {
return zero, nil, errors.New("tool handler must have 2 arguments")
}
if handlerType.NumOut() != 2 {
return zero, nil, errors.New("tool handler must return 2 values")
}
if handlerType.In(0) != reflect.TypeOf((*context.Context)(nil)).Elem() {
return zero, nil, errors.New("tool handler first argument must be context.Context")
}
// We no longer check the type of the first return value
if handlerType.Out(1).Kind() != reflect.Interface {
return zero, nil, errors.New("tool handler second return value must be error")
}
argType := handlerType.In(1)
if argType.Kind() != reflect.Struct {
return zero, nil, errors.New("tool handler second argument must be a struct")
}
handler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
// Create OpenTelemetry span for tool execution (no-op when no exporter configured)
config := GrafanaConfigFromContext(ctx)
ctx, span := otel.Tracer("mcp-grafana").Start(ctx, fmt.Sprintf("mcp.tool.%s", name))
defer span.End()
// Add tool metadata as span attributes
span.SetAttributes(
attribute.String("mcp.tool.name", name),
attribute.String("mcp.tool.description", description),
)
argBytes, err := json.Marshal(request.Params.Arguments)
if err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, "failed to marshal arguments")
return nil, fmt.Errorf("marshal args: %w", err)
}
// Add arguments as span attribute only if adding args to trace attributes is enabled
if config.IncludeArgumentsInSpans {
span.SetAttributes(attribute.String("mcp.tool.arguments", string(argBytes)))
}
unmarshaledArgs := reflect.New(argType).Interface()
if err := json.Unmarshal(argBytes, unmarshaledArgs); err != nil {
span.RecordError(err)
span.SetStatus(codes.Error, "failed to unmarshal arguments")
return nil, fmt.Errorf("unmarshal args: %s", err)
}
// Need to dereference the unmarshaled arguments
of := reflect.ValueOf(unmarshaledArgs)
if of.Kind() != reflect.Ptr || !of.Elem().CanInterface() {
err := errors.New("arguments must be a struct")
span.RecordError(err)
span.SetStatus(codes.Error, "invalid arguments structure")
return nil, err
}
// Pass the instrumented context to the tool handler
args := []reflect.Value{reflect.ValueOf(ctx), of.Elem()}
output := handlerValue.Call(args)
if len(output) != 2 {
err := errors.New("tool handler must return 2 values")
span.RecordError(err)
span.SetStatus(codes.Error, "invalid tool handler return")
return nil, err
}
if !output[0].CanInterface() {
err := errors.New("tool handler first return value must be interfaceable")
span.RecordError(err)
span.SetStatus(codes.Error, "tool handler return value not interfaceable")
return nil, err
}
// Handle the error return value first
var handlerErr error
var ok bool
if output[1].Kind() == reflect.Interface && !output[1].IsNil() {
handlerErr, ok = output[1].Interface().(error)
if !ok {
err := errors.New("tool handler second return value must be error")
span.RecordError(err)
span.SetStatus(codes.Error, "invalid error return type")
return nil, err
}
}
// If there's an error, record it and return
if handlerErr != nil {
span.RecordError(handlerErr)
span.SetStatus(codes.Error, handlerErr.Error())
return nil, handlerErr
}
// Tool execution completed successfully
span.SetStatus(codes.Ok, "tool execution completed")
// Check if the first return value is nil (only for pointer, interface, map, etc.)
isNilable := output[0].Kind() == reflect.Ptr ||
output[0].Kind() == reflect.Interface ||
output[0].Kind() == reflect.Map ||
output[0].Kind() == reflect.Slice ||
output[0].Kind() == reflect.Chan ||
output[0].Kind() == reflect.Func
if isNilable && output[0].IsNil() {
return nil, nil
}
returnVal := output[0].Interface()
returnType := output[0].Type()
// Case 1: Already a *mcp.CallToolResult
if callResult, ok := returnVal.(*mcp.CallToolResult); ok {
return callResult, nil
}
// Case 2: An mcp.CallToolResult (not a pointer)
if returnType.ConvertibleTo(reflect.TypeOf(mcp.CallToolResult{})) {
callResult := returnVal.(mcp.CallToolResult)
return &callResult, nil
}
// Case 3: String or *string
if str, ok := returnVal.(string); ok {
if str == "" {
return nil, nil
}
return mcp.NewToolResultText(str), nil
}
if strPtr, ok := returnVal.(*string); ok {
if strPtr == nil || *strPtr == "" {
return nil, nil
}
return mcp.NewToolResultText(*strPtr), nil
}
// Case 4: Any other type - marshal to JSON
returnBytes, err := json.Marshal(returnVal)
if err != nil {
return nil, fmt.Errorf("failed to marshal return value: %s", err)
}
return mcp.NewToolResultText(string(returnBytes)), nil
}
jsonSchema := createJSONSchemaFromHandler(toolHandler)
properties := make(map[string]any, jsonSchema.Properties.Len())
for pair := jsonSchema.Properties.Oldest(); pair != nil; pair = pair.Next() {
properties[pair.Key] = pair.Value
}
inputSchema := mcp.ToolInputSchema{
Type: jsonSchema.Type,
Properties: properties,
Required: jsonSchema.Required,
}
t := mcp.Tool{
Name: name,
Description: description,
InputSchema: inputSchema,
}
for _, option := range options {
option(&t)
}
return t, handler, nil
}
// Creates a full JSON schema from a user provided handler by introspecting the arguments
func createJSONSchemaFromHandler(handler any) *jsonschema.Schema {
handlerValue := reflect.ValueOf(handler)
handlerType := handlerValue.Type()
argumentType := handlerType.In(1)
inputSchema := jsonSchemaReflector.ReflectFromType(argumentType)
return inputSchema
}
var (
jsonSchemaReflector = jsonschema.Reflector{
BaseSchemaID: "",
Anonymous: true,
AssignAnchor: false,
AllowAdditionalProperties: true,
RequiredFromJSONSchemaTags: true,
DoNotReference: true,
ExpandedStruct: true,
FieldNameTag: "",
IgnoredTypes: nil,
Lookup: nil,
Mapper: nil,
Namer: nil,
KeyNamer: nil,
AdditionalFields: nil,
CommentMap: nil,
}
)
```
--------------------------------------------------------------------------------
/tools/prometheus_test.go:
--------------------------------------------------------------------------------
```go
// Requires a Grafana instance running on localhost:3000,
// with a Prometheus datasource provisioned.
// Run with `go test -tags integration`.
//go:build integration
package tools
import (
"fmt"
"testing"
"time"
"github.com/prometheus/common/model"
"github.com/prometheus/prometheus/model/labels"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestPrometheusTools(t *testing.T) {
t.Run("list prometheus metric metadata", func(t *testing.T) {
ctx := newTestContext()
result, err := listPrometheusMetricMetadata(ctx, ListPrometheusMetricMetadataParams{
DatasourceUID: "prometheus",
})
require.NoError(t, err)
assert.Len(t, result, 10)
})
t.Run("list prometheus metric names", func(t *testing.T) {
ctx := newTestContext()
result, err := listPrometheusMetricNames(ctx, ListPrometheusMetricNamesParams{
DatasourceUID: "prometheus",
Regex: ".*",
Limit: 10,
})
require.NoError(t, err)
assert.Len(t, result, 10)
})
t.Run("list prometheus label names", func(t *testing.T) {
ctx := newTestContext()
result, err := listPrometheusLabelNames(ctx, ListPrometheusLabelNamesParams{
DatasourceUID: "prometheus",
Matches: []Selector{
{
Filters: []LabelMatcher{
{Name: "job", Value: "prometheus"},
},
},
},
Limit: 10,
})
require.NoError(t, err)
assert.Len(t, result, 10)
})
t.Run("list prometheus label values", func(t *testing.T) {
ctx := newTestContext()
result, err := listPrometheusLabelValues(ctx, ListPrometheusLabelValuesParams{
DatasourceUID: "prometheus",
LabelName: "job",
Matches: []Selector{
{
Filters: []LabelMatcher{
{Name: "job", Value: "prometheus"},
},
},
},
})
require.NoError(t, err)
assert.Len(t, result, 1)
})
}
func TestSelectorMatches(t *testing.T) {
testCases := []struct {
name string
selector Selector
labels map[string]string
expected bool
expectErr bool
}{
{
name: "Equal match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=", Value: "prometheus"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: true,
},
{
name: "Equal no match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=", Value: "prometheus"},
},
},
labels: map[string]string{"job": "node-exporter"},
expected: false,
},
{
name: "Not equal match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "!=", Value: "prometheus"},
},
},
labels: map[string]string{"job": "node-exporter"},
expected: true,
},
{
name: "Not equal no match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "!=", Value: "prometheus"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: false,
},
{
name: "Regex match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=~", Value: "prom.*"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: true,
},
{
name: "Regex no match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=~", Value: "node.*"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: false,
},
{
name: "Not regex match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "!~", Value: "node.*"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: true,
},
{
name: "Not regex no match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "!~", Value: "prom.*"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: false,
},
{
name: "Multiple filters all match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=", Value: "prometheus"},
{Name: "instance", Type: "=~", Value: "localhost.*"},
},
},
labels: map[string]string{"job": "prometheus", "instance": "localhost:9090"},
expected: true,
},
{
name: "Multiple filters one doesn't match",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "=", Value: "prometheus"},
{Name: "instance", Type: "=~", Value: "remote.*"},
},
},
labels: map[string]string{"job": "prometheus", "instance": "localhost:9090"},
expected: false,
},
{
name: "Label doesn't exist with = operator",
selector: Selector{
Filters: []LabelMatcher{
{Name: "missing", Type: "=", Value: "value"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: false,
},
{
name: "Label doesn't exist with != operator",
selector: Selector{
Filters: []LabelMatcher{
{Name: "missing", Type: "!=", Value: "value"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: true,
},
{
name: "Invalid matcher type",
selector: Selector{
Filters: []LabelMatcher{
{Name: "job", Type: "<>", Value: "prometheus"},
},
},
labels: map[string]string{"job": "prometheus"},
expected: false,
expectErr: true,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
lbls := labels.FromMap(tc.labels)
result, err := tc.selector.Matches(lbls)
if tc.expectErr {
assert.Error(t, err)
return
}
assert.NoError(t, err)
assert.Equal(t, tc.expected, result)
})
}
}
func TestPrometheusQueries(t *testing.T) {
t.Run("query prometheus range", func(t *testing.T) {
end := time.Now()
start := end.Add(-10 * time.Minute)
for _, step := range []int{15, 60, 300} {
t.Run(fmt.Sprintf("step=%d", step), func(t *testing.T) {
ctx := newTestContext()
result, err := queryPrometheus(ctx, QueryPrometheusParams{
DatasourceUID: "prometheus",
Expr: "test",
StartTime: start.Format(time.RFC3339),
EndTime: end.Format(time.RFC3339),
StepSeconds: step,
QueryType: "range",
})
require.NoError(t, err)
matrix := result.(model.Matrix)
require.Len(t, matrix, 1)
expectedLen := int(end.Sub(start).Seconds()/float64(step)) + 1
assert.Len(t, matrix[0].Values, expectedLen)
assert.Less(t, matrix[0].Values[0].Timestamp.Sub(model.TimeFromUnix(start.Unix())), time.Duration(step)*time.Second)
assert.Equal(t, matrix[0].Metric["__name__"], model.LabelValue("test"))
})
}
})
t.Run("query prometheus instant", func(t *testing.T) {
ctx := newTestContext()
result, err := queryPrometheus(ctx, QueryPrometheusParams{
DatasourceUID: "prometheus",
Expr: "up",
StartTime: time.Now().Format(time.RFC3339),
QueryType: "instant",
})
require.NoError(t, err)
scalar := result.(model.Vector)
assert.Equal(t, scalar[0].Value, model.SampleValue(1))
assert.Equal(t, scalar[0].Timestamp, model.TimeFromUnix(time.Now().Unix()))
assert.Equal(t, scalar[0].Metric["__name__"], model.LabelValue("up"))
})
t.Run("query prometheus instant with relative timestamps", func(t *testing.T) {
ctx := newTestContext()
beforeQuery := model.TimeFromUnix(time.Now().Unix())
result, err := queryPrometheus(ctx, QueryPrometheusParams{
DatasourceUID: "prometheus",
Expr: "up",
StartTime: "now",
QueryType: "instant",
})
afterQuery := model.TimeFromUnix(time.Now().Unix())
require.NoError(t, err)
scalar := result.(model.Vector)
assert.Equal(t, scalar[0].Value, model.SampleValue(1))
// Check that the timestamp is within the expected range
buffer := 5 * time.Second
assert.True(t, scalar[0].Timestamp >= beforeQuery,
"Result timestamp should be after or equal to the time before the query")
assert.True(t, scalar[0].Timestamp <= afterQuery.Add(buffer),
"Result timestamp should be before or equal to the time after the query (with 5s buffer)")
assert.Equal(t, scalar[0].Metric["__name__"], model.LabelValue("up"))
})
t.Run("query prometheus range with relative timestamps", func(t *testing.T) {
ctx := newTestContext()
beforeQuery := model.TimeFromUnix(time.Now().Unix())
result, err := queryPrometheus(ctx, QueryPrometheusParams{
DatasourceUID: "prometheus",
Expr: "test",
StartTime: "now-1h",
EndTime: "now",
StepSeconds: 60,
QueryType: "range",
})
afterQuery := model.TimeFromUnix(time.Now().Unix())
require.NoError(t, err)
matrix := result.(model.Matrix)
require.Len(t, matrix, 1)
// Should have approximately 60 samples (one per minute for an hour)
assert.InDelta(t, 60, len(matrix[0].Values), 2)
buffer := 5 * time.Second
oneHour := time.Hour
firstSampleTime := matrix[0].Values[0].Timestamp
// Check that the start timestamp is within the expected range
assert.True(t, firstSampleTime >= beforeQuery.Add(-oneHour),
"First timestamp should be after or equal to the time before the query minus one hour")
assert.True(t, firstSampleTime <= afterQuery.Add(buffer).Add(-oneHour),
"First timestamp should be before or equal to the time after the query minus one hour (with 5s buffer)")
// Check that the end timestamp is is within the expected range
lastSampleTime := matrix[0].Values[len(matrix[0].Values)-1].Timestamp
assert.True(t, lastSampleTime >= beforeQuery,
"Last timestamp should be after or equal to the time before the query")
assert.True(t, lastSampleTime <= afterQuery.Add(buffer),
"Last timestamp should be before or equal to the time after the query (with 5s buffer)")
assert.Equal(t, matrix[0].Metric["__name__"], model.LabelValue("test"))
})
}
```
--------------------------------------------------------------------------------
/tests/tempo_test.py:
--------------------------------------------------------------------------------
```python
from mcp import ClientSession
import pytest
from langevals import expect
from langevals_langevals.llm_boolean import (
CustomLLMBooleanEvaluator,
CustomLLMBooleanSettings,
)
from litellm import Message, acompletion
from mcp import ClientSession
from conftest import models
from utils import (
get_converted_tools,
llm_tool_call_sequence,
)
pytestmark = pytest.mark.anyio
class TestTempoProxiedToolsBasic:
"""Test Tempo proxied MCP tools functionality.
These tests verify that Tempo datasources with MCP support are discovered
per-session and their tools are registered with a datasourceUid parameter
for multi-datasource support.
Requires:
- Docker compose services running (includes 2 Tempo instances)
- GRAFANA_USERNAME and GRAFANA_PASSWORD environment variables
- MCP server running
"""
@pytest.mark.anyio
async def test_tempo_tools_discovered_and_registered(
self, mcp_client: ClientSession
):
"""Test that Tempo tools are discovered and registered with datasourceUid parameter."""
# List all tools
list_response = await mcp_client.list_tools()
all_tool_names = [tool.name for tool in list_response.tools]
# Find tempo-prefixed tools (should preserve hyphens from original tool names)
tempo_tools = [name for name in all_tool_names if name.startswith("tempo_")]
# Expected tools from Tempo MCP server
expected_tempo_tools = [
"tempo_traceql-search",
"tempo_traceql-metrics-instant",
"tempo_traceql-metrics-range",
"tempo_get-trace",
"tempo_get-attribute-names",
"tempo_get-attribute-values",
"tempo_docs-traceql",
]
assert len(tempo_tools) == len(expected_tempo_tools), (
f"Expected {len(expected_tempo_tools)} unique tempo tools, found {len(tempo_tools)}: {tempo_tools}"
)
for expected_tool in expected_tempo_tools:
assert expected_tool in tempo_tools, (
f"Tool {expected_tool} should be available"
)
@pytest.mark.anyio
async def test_tempo_tools_have_datasourceUid_parameter(self, mcp_client):
"""Test that all tempo tools have a required datasourceUid parameter."""
list_response = await mcp_client.list_tools()
tempo_tools = [
tool for tool in list_response.tools if tool.name.startswith("tempo_")
]
assert len(tempo_tools) > 0, "Should have at least one tempo tool"
for tool in tempo_tools:
# Verify the tool has input schema
assert hasattr(tool, "inputSchema"), (
f"Tool {tool.name} should have inputSchema"
)
assert isinstance(tool.inputSchema, dict), (
f"Tool {tool.name} inputSchema should be a dict"
)
# Verify datasourceUid parameter exists (camelCase)
properties = tool.inputSchema.get("properties", {})
assert "datasourceUid" in properties, (
f"Tool {tool.name} should have datasourceUid parameter (camelCase)"
)
# Verify it's required
required = tool.inputSchema.get("required", [])
assert "datasourceUid" in required, (
f"Tool {tool.name} should require datasourceUid parameter"
)
# Verify parameter has proper description
datasource_uid_prop = properties["datasourceUid"]
assert "type" in datasource_uid_prop, (
f"datasourceUid should have type defined"
)
assert datasource_uid_prop["type"] == "string", (
f"datasourceUid should be type string"
)
@pytest.mark.anyio
async def test_tempo_tool_call_with_valid_datasource(self, mcp_client):
"""Test calling a tempo tool with a valid datasourceUid."""
# Call docs-traceql which should return documentation (doesn't require data)
try:
call_response = await mcp_client.call_tool(
"tempo_docs-traceql",
arguments={"datasourceUid": "tempo", "name": "basic"},
)
# Verify we got a response
assert call_response.content, "Tool should return content"
# Should have text content (documentation)
response_text = call_response.content[0].text
assert len(response_text) > 0, "Response should have content"
assert "traceql" in response_text.lower(), (
"Response should contain TraceQL documentation"
)
print(response_text)
except Exception as e:
# If this fails, it might be because Tempo doesn't have data yet
# but at least verify the error isn't about missing datasourceUid
error_msg = str(e).lower()
assert "datasourceuid" not in error_msg, (
f"Should not fail due to datasourceUid parameter: {e}"
)
print(error_msg)
@pytest.mark.anyio
async def test_tempo_tool_call_missing_datasourceUid(self, mcp_client):
"""Test that calling a tempo tool without datasourceUid fails appropriately."""
with pytest.raises(Exception) as exc_info:
await mcp_client.call_tool(
"tempo_docs-traceql",
arguments={"name": "basic"}, # Missing datasourceUid
)
error_msg = str(exc_info.value).lower()
assert "datasourceuid" in error_msg and "required" in error_msg, (
f"Should require datasourceUid parameter: {exc_info.value}"
)
@pytest.mark.anyio
async def test_tempo_tool_call_invalid_datasourceUid(self, mcp_client):
"""Test that calling a tempo tool with invalid datasourceUid returns helpful error."""
with pytest.raises(Exception) as exc_info:
await mcp_client.call_tool(
"tempo_docs-traceql",
arguments={"datasourceUid": "nonexistent-tempo", "name": "basic"},
)
error_msg = str(exc_info.value).lower()
# Should mention that datasource wasn't found
assert "not found" in error_msg or "not accessible" in error_msg, (
f"Should indicate datasource not found: {exc_info.value}"
)
# Should mention available datasources to help user
assert "tempo" in error_msg or "available" in error_msg, (
f"Error should be helpful and mention available datasources: {exc_info.value}"
)
@pytest.mark.anyio
async def test_tempo_tool_works_with_multiple_datasources(self, mcp_client):
"""Test that the same tool works with different datasources via datasourceUid."""
# Both tempo and tempo-secondary should be available in our test environment
datasources = ["tempo", "tempo-secondary"]
for datasource_uid in datasources:
try:
# Call the same tool with different datasources
call_response = await mcp_client.call_tool(
"tempo_get-attribute-names",
arguments={"datasourceUid": datasource_uid},
)
# Verify we got a response
assert call_response.content, (
f"Tool should return content for datasource {datasource_uid}"
)
# Response should be valid JSON or text
response_text = call_response.content[0].text
assert len(response_text) > 0, (
f"Response should have content for datasource {datasource_uid}"
)
except Exception as e:
# If this fails, it's acceptable if Tempo doesn't have trace data yet
# But verify it's not a routing/config error
error_msg = str(e).lower()
assert (
"not found" not in error_msg or datasource_uid not in error_msg
), f"Datasource {datasource_uid} should be accessible: {e}"
class TestTempoProxiedToolsWithLLM:
"""LLM integration tests for Tempo proxied tools."""
@pytest.mark.parametrize("model", models)
@pytest.mark.flaky(max_runs=3)
async def test_llm_can_list_trace_attributes(
self, model: str, mcp_client: ClientSession
):
"""Test that an LLM can list available trace attributes from Tempo."""
tools = await get_converted_tools(mcp_client)
prompt = (
"Use the tempo tools to get a list of all available trace attribute names "
"from the datasource with UID 'tempo'. I want to know what attributes "
"I can use in my TraceQL queries."
)
messages = [
Message(role="system", content="You are a helpful assistant."),
Message(role="user", content=prompt),
]
# LLM should call tempo_get-attribute-names with datasourceUid
messages = await llm_tool_call_sequence(
model,
messages,
tools,
mcp_client,
"tempo_get-attribute-names",
{"datasourceUid": "tempo"},
)
# Final LLM response should mention attributes
response = await acompletion(model=model, messages=messages, tools=tools)
content = response.choices[0].message.content
attributes_checker = CustomLLMBooleanEvaluator(
settings=CustomLLMBooleanSettings(
prompt="Does the response list or describe trace attributes that are available for querying?",
)
)
expect(input=prompt, output=content).to_pass(attributes_checker)
```
--------------------------------------------------------------------------------
/tools/annotations.go:
--------------------------------------------------------------------------------
```go
package tools
import (
"context"
"fmt"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
"strconv"
mcpgrafana "github.com/grafana/mcp-grafana"
"github.com/grafana/grafana-openapi-client-go/client/annotations"
"github.com/grafana/grafana-openapi-client-go/models"
)
// GetAnnotationsInput filters annotation search.
type GetAnnotationsInput struct {
From *int64 `jsonschema:"description=Epoch ms start time"`
To *int64 `jsonschema:"description=Epoch ms end time"`
Limit *int64 `jsonschema:"description=Max results default 100"`
AlertID *int64 `jsonschema:"description=Deprecated. Use AlertUID"`
AlertUID *string `jsonschema:"description=Filter by alert UID"`
DashboardID *int64 `jsonschema:"description=Deprecated. Use DashboardUID"`
DashboardUID *string `jsonschema:"description=Filter by dashboard UID"`
PanelID *int64 `jsonschema:"description=Filter by panel ID"`
UserID *int64 `jsonschema:"description=Filter by creator user ID"`
Type *string `jsonschema:"description=annotation or alert"`
Tags []string `jsonschema:"description=Multiple tags allowed tags=tag1&tags=tag2"`
MatchAny *bool `jsonschema:"description=true OR tag match false AND"`
}
// getAnnotations retrieves Grafana annotations using filters.
func getAnnotations(ctx context.Context, args GetAnnotationsInput) (*annotations.GetAnnotationsOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
req := annotations.GetAnnotationsParams{
From: args.From,
To: args.To,
Limit: args.Limit,
AlertID: args.AlertID,
AlertUID: args.AlertUID,
DashboardID: args.DashboardID,
DashboardUID: args.DashboardUID,
PanelID: args.PanelID,
UserID: args.UserID,
Type: args.Type,
Tags: args.Tags,
MatchAny: args.MatchAny,
Context: ctx,
}
resp, err := c.Annotations.GetAnnotations(&req)
if err != nil {
return nil, fmt.Errorf("get annotations: %w", err)
}
return resp, nil
}
var GetAnnotationsTool = mcpgrafana.MustTool(
"get_annotations",
"Fetch Grafana annotations using filters such as dashboard UID, time range and tags.",
getAnnotations,
mcp.WithTitleAnnotation("Get Annotations"),
mcp.WithIdempotentHintAnnotation(true),
mcp.WithReadOnlyHintAnnotation(true),
)
// CreateAnnotationInput creates a new annotation.
type CreateAnnotationInput struct {
DashboardID int64 `json:"dashboardId,omitempty" jsonschema:"description=Deprecated. Use dashboardUID"`
DashboardUID string `json:"dashboardUID,omitempty" jsonschema:"description=Preferred dashboard UID"`
PanelID int64 `json:"panelId,omitempty" jsonschema:"description=Panel ID"`
Time int64 `json:"time,omitempty" jsonschema:"description=Start time epoch ms"`
TimeEnd int64 `json:"timeEnd,omitempty" jsonschema:"description=End time epoch ms"`
Tags []string `json:"tags,omitempty" jsonschema:"description=Optional list of tags"`
Text string `json:"text" jsonschema:"description=Annotation text required"`
Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional JSON payload"`
}
// createAnnotation sends a POST request to create a Grafana annotation.
func createAnnotation(ctx context.Context, args CreateAnnotationInput) (*annotations.PostAnnotationOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
req := models.PostAnnotationsCmd{
DashboardID: args.DashboardID,
DashboardUID: args.DashboardUID,
PanelID: args.PanelID,
Time: args.Time,
TimeEnd: args.TimeEnd,
Tags: args.Tags,
Text: &args.Text,
Data: args.Data,
}
resp, err := c.Annotations.PostAnnotation(&req)
if err != nil {
return nil, fmt.Errorf("create annotation: %w", err)
}
return resp, nil
}
var CreateAnnotationTool = mcpgrafana.MustTool(
"create_annotation",
"Create a new annotation on a dashboard or panel.",
createAnnotation,
mcp.WithTitleAnnotation("Create Annotation"),
mcp.WithIdempotentHintAnnotation(false),
)
// CreateGraphiteAnnotationInput represents the payload format for creating a Graphite-style annotation.
type CreateGraphiteAnnotationInput struct {
What string `json:"what" jsonschema:"description=Annotation text"`
When int64 `json:"when" jsonschema:"description=Epoch ms timestamp"`
Tags []string `json:"tags,omitempty" jsonschema:"description=Optional list of tags"`
Data string `json:"data,omitempty" jsonschema:"description=Optional payload"`
}
// createAnnotationGraphiteFormat creates an annotation using the Graphite annotation format.
func createAnnotationGraphiteFormat(ctx context.Context, args CreateGraphiteAnnotationInput) (*annotations.PostGraphiteAnnotationOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
req := &models.PostGraphiteAnnotationsCmd{
What: args.What,
When: args.When,
Tags: args.Tags,
Data: args.Data,
}
resp, err := c.Annotations.PostGraphiteAnnotation(req)
if err != nil {
return nil, fmt.Errorf("create graphite annotation: %w", err)
}
return resp, nil
}
var CreateGraphiteAnnotationTool = mcpgrafana.MustTool(
"create_graphite_annotation",
"Create an annotation using Graphite annotation format.",
createAnnotationGraphiteFormat,
mcp.WithTitleAnnotation("Create Graphite Annotation"),
mcp.WithIdempotentHintAnnotation(false),
)
// UpdateAnnotationInput represents the payload used to update an existing annotation by ID.
type UpdateAnnotationInput struct {
ID int64 `json:"id" jsonschema:"description=Annotation ID to update"`
Time int64 `json:"time,omitempty" jsonschema:"description=Start time epoch ms"`
TimeEnd int64 `json:"timeEnd,omitempty" jsonschema:"description=End time epoch ms"`
Text string `json:"text,omitempty" jsonschema:"description=Annotation text"`
Tags []string `json:"tags,omitempty" jsonschema:"description=Tags to replace existing tags"`
Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional JSON payload"`
}
// updateAnnotation updates an annotation using its ID.
func updateAnnotation(ctx context.Context, args UpdateAnnotationInput) (*annotations.UpdateAnnotationOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
annotationID := strconv.FormatInt(args.ID, 10)
req := &models.UpdateAnnotationsCmd{
Time: args.Time,
TimeEnd: args.TimeEnd,
Text: args.Text,
Tags: args.Tags,
Data: args.Data,
}
resp, err := c.Annotations.UpdateAnnotation(annotationID, req)
if err != nil {
return nil, fmt.Errorf("update annotation: %w", err)
}
return resp, nil
}
var UpdateAnnotationTool = mcpgrafana.MustTool(
"update_annotation",
"Updates all properties of an annotation that matches the specified ID. Sends a full update (PUT). For partial updates, use patch_annotation instead.",
updateAnnotation,
mcp.WithTitleAnnotation("Update Annotation"),
mcp.WithIdempotentHintAnnotation(false),
)
// PatchAnnotationInput updates only the provided fields.
type PatchAnnotationInput struct {
ID int64 `json:"id" jsonschema:"description=Annotation ID"`
Text *string `json:"text,omitempty" jsonschema:"description=Optional new text"`
Tags []string `json:"tags,omitempty" jsonschema:"description=Optional replace tags"`
Time *int64 `json:"time,omitempty" jsonschema:"description=Optional new start epoch ms"`
TimeEnd *int64 `json:"timeEnd,omitempty" jsonschema:"description=Optional new end epoch ms"`
Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional metadata"`
}
// patchAnnotation patches only the provided annotation fields.
func patchAnnotation(ctx context.Context, args PatchAnnotationInput) (*annotations.PatchAnnotationOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
id := strconv.FormatInt(args.ID, 10)
body := &models.PatchAnnotationsCmd{}
if args.Text != nil {
body.Text = *args.Text
}
if args.Time != nil {
body.Time = *args.Time
}
if args.TimeEnd != nil {
body.TimeEnd = *args.TimeEnd
}
if args.Tags != nil {
body.Tags = args.Tags
}
if args.Data != nil {
body.Data = args.Data
}
resp, err := c.Annotations.PatchAnnotation(id, body)
if err != nil {
return nil, fmt.Errorf("patch annotation: %w", err)
}
return resp, nil
}
var PatchAnnotationTool = mcpgrafana.MustTool(
"patch_annotation",
"Updates only the provided properties of an annotation. Fields omitted are not modified. Use update_annotation for full replacement.",
patchAnnotation,
mcp.WithTitleAnnotation("Patch Annotation"),
mcp.WithIdempotentHintAnnotation(false),
)
// GetAnnotationTagsInput defines filters for retrieving annotation tags.
type GetAnnotationTagsInput struct {
Tag *string `json:"tag,omitempty" jsonschema:"description=Optional filter by tag name"`
Limit *string `json:"limit,omitempty" jsonschema:"description=Max results\\, default 100"`
}
func getAnnotationTags(ctx context.Context, args GetAnnotationTagsInput) (*annotations.GetAnnotationTagsOK, error) {
c := mcpgrafana.GrafanaClientFromContext(ctx)
req := annotations.GetAnnotationTagsParams{
Tag: args.Tag,
Limit: args.Limit,
Context: ctx,
}
resp, err := c.Annotations.GetAnnotationTags(&req)
if err != nil {
return nil, fmt.Errorf("get annotation tags: %w", err)
}
return resp, nil
}
var GetAnnotationTagsTool = mcpgrafana.MustTool(
"get_annotation_tags",
"Returns annotation tags with optional filtering by tag name. Only the provided filters are applied.",
getAnnotationTags,
mcp.WithTitleAnnotation("Get Annotation Tags"),
mcp.WithIdempotentHintAnnotation(true),
mcp.WithReadOnlyHintAnnotation(true),
)
func AddAnnotationTools(mcp *server.MCPServer, enableWriteTools bool) {
GetAnnotationsTool.Register(mcp)
if enableWriteTools {
CreateAnnotationTool.Register(mcp)
CreateGraphiteAnnotationTool.Register(mcp)
UpdateAnnotationTool.Register(mcp)
PatchAnnotationTool.Register(mcp)
}
GetAnnotationTagsTool.Register(mcp)
}
```
--------------------------------------------------------------------------------
/tools/oncall_cloud_test.go:
--------------------------------------------------------------------------------
```go
//go:build cloud
// +build cloud
// This file contains cloud integration tests that run against a dedicated test instance
// at mcptests.grafana-dev.net. This instance is configured with a minimal setup on the OnCall side:
// - One team
// - Two schedules (only one has a team assigned)
// - One shift in the schedule with a team assigned
// - One user
// These tests expect this configuration to exist and will skip if the required
// environment variables (GRAFANA_URL, GRAFANA_SERVICE_ACCOUNT_TOKEN or GRAFANA_API_KEY) are not set.
// The GRAFANA_API_KEY variable is deprecated.
package tools
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestCloudOnCallSchedules(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
// Test listing all schedules
t.Run("list all schedules", func(t *testing.T) {
result, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{})
require.NoError(t, err, "Should not error when listing schedules")
assert.NotNil(t, result, "Result should not be nil")
})
// Test pagination
t.Run("list schedules with pagination", func(t *testing.T) {
// Get first page
page1, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{Page: 1})
require.NoError(t, err, "Should not error when listing schedules page 1")
assert.NotNil(t, page1, "Page 1 should not be nil")
// Get second page
page2, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{Page: 2})
require.NoError(t, err, "Should not error when listing schedules page 2")
assert.NotNil(t, page2, "Page 2 should not be nil")
})
// Get a team ID from an existing schedule to test filtering
schedules, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{})
require.NoError(t, err, "Should not error when listing schedules")
if len(schedules) > 0 && schedules[0].TeamID != "" {
teamID := schedules[0].TeamID
// Test filtering by team ID
t.Run("list schedules by team ID", func(t *testing.T) {
result, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{
TeamID: teamID,
})
require.NoError(t, err, "Should not error when listing schedules by team")
assert.NotEmpty(t, result, "Should return at least one schedule")
for _, schedule := range result {
assert.Equal(t, teamID, schedule.TeamID, "All schedules should belong to the specified team")
}
})
}
// Test getting a specific schedule
if len(schedules) > 0 {
scheduleID := schedules[0].ID
t.Run("get specific schedule", func(t *testing.T) {
result, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{
ScheduleID: scheduleID,
})
require.NoError(t, err, "Should not error when getting specific schedule")
assert.Len(t, result, 1, "Should return exactly one schedule")
assert.Equal(t, scheduleID, result[0].ID, "Should return the correct schedule")
// Verify all summary fields are present
schedule := result[0]
assert.NotEmpty(t, schedule.Name, "Schedule should have a name")
assert.NotEmpty(t, schedule.Timezone, "Schedule should have a timezone")
assert.NotNil(t, schedule.Shifts, "Schedule should have a shifts field")
})
}
}
func TestCloudOnCallShift(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
// First get a schedule to find a valid shift
schedules, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{})
require.NoError(t, err, "Should not error when listing schedules")
require.NotEmpty(t, schedules, "Should have at least one schedule to test with")
require.NotEmpty(t, schedules[0].Shifts, "Schedule should have at least one shift")
shifts := schedules[0].Shifts
shiftID := shifts[0]
// Test getting shift details with valid ID
t.Run("get shift details", func(t *testing.T) {
result, err := getOnCallShift(ctx, GetOnCallShiftParams{
ShiftID: shiftID,
})
require.NoError(t, err, "Should not error when getting shift details")
assert.NotNil(t, result, "Result should not be nil")
assert.Equal(t, shiftID, result.ID, "Should return the correct shift")
})
t.Run("get shift with invalid ID", func(t *testing.T) {
_, err := getOnCallShift(ctx, GetOnCallShiftParams{
ShiftID: "invalid-shift-id",
})
assert.Error(t, err, "Should error when getting shift with invalid ID")
})
}
func TestCloudGetCurrentOnCallUsers(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
// First get a schedule to use for testing
schedules, err := listOnCallSchedules(ctx, ListOnCallSchedulesParams{})
require.NoError(t, err, "Should not error when listing schedules")
require.NotEmpty(t, schedules, "Should have at least one schedule to test with")
scheduleID := schedules[0].ID
// Test getting current on-call users
t.Run("get current on-call users", func(t *testing.T) {
result, err := getCurrentOnCallUsers(ctx, GetCurrentOnCallUsersParams{
ScheduleID: scheduleID,
})
require.NoError(t, err, "Should not error when getting current on-call users")
assert.NotNil(t, result, "Result should not be nil")
assert.Equal(t, scheduleID, result.ScheduleID, "Should return the correct schedule")
assert.NotEmpty(t, result.ScheduleName, "Schedule should have a name")
assert.NotNil(t, result.Users, "Users field should be present")
// Assert that Users is of type []*aapi.User
if len(result.Users) > 0 {
user := result.Users[0]
assert.NotEmpty(t, user.ID, "User should have an ID")
assert.NotEmpty(t, user.Username, "User should have a username")
}
})
t.Run("get current on-call users with invalid schedule ID", func(t *testing.T) {
_, err := getCurrentOnCallUsers(ctx, GetCurrentOnCallUsersParams{
ScheduleID: "invalid-schedule-id",
})
assert.Error(t, err, "Should error when getting current on-call users with invalid schedule ID")
})
}
func TestCloudOnCallTeams(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
t.Run("list teams", func(t *testing.T) {
result, err := listOnCallTeams(ctx, ListOnCallTeamsParams{})
require.NoError(t, err, "Should not error when listing teams")
assert.NotNil(t, result, "Result should not be nil")
if len(result) > 0 {
team := result[0]
assert.NotEmpty(t, team.ID, "Team should have an ID")
assert.NotEmpty(t, team.Name, "Team should have a name")
}
})
// Test pagination
t.Run("list teams with pagination", func(t *testing.T) {
// Get first page
page1, err := listOnCallTeams(ctx, ListOnCallTeamsParams{Page: 1})
require.NoError(t, err, "Should not error when listing teams page 1")
assert.NotNil(t, page1, "Page 1 should not be nil")
// Get second page
page2, err := listOnCallTeams(ctx, ListOnCallTeamsParams{Page: 2})
require.NoError(t, err, "Should not error when listing teams page 2")
assert.NotNil(t, page2, "Page 2 should not be nil")
})
}
func TestCloudOnCallUsers(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
t.Run("list all users", func(t *testing.T) {
result, err := listOnCallUsers(ctx, ListOnCallUsersParams{})
require.NoError(t, err, "Should not error when listing users")
assert.NotNil(t, result, "Result should not be nil")
if len(result) > 0 {
user := result[0]
assert.NotEmpty(t, user.ID, "User should have an ID")
assert.NotEmpty(t, user.Username, "User should have a username")
}
})
// Test pagination
t.Run("list users with pagination", func(t *testing.T) {
// Get first page
page1, err := listOnCallUsers(ctx, ListOnCallUsersParams{Page: 1})
require.NoError(t, err, "Should not error when listing users page 1")
assert.NotNil(t, page1, "Page 1 should not be nil")
// Get second page
page2, err := listOnCallUsers(ctx, ListOnCallUsersParams{Page: 2})
require.NoError(t, err, "Should not error when listing users page 2")
assert.NotNil(t, page2, "Page 2 should not be nil")
})
// Get a user ID and username from the list to test filtering
users, err := listOnCallUsers(ctx, ListOnCallUsersParams{})
require.NoError(t, err, "Should not error when listing users")
require.NotEmpty(t, users, "Should have at least one user to test with")
userID := users[0].ID
username := users[0].Username
t.Run("get user by ID", func(t *testing.T) {
result, err := listOnCallUsers(ctx, ListOnCallUsersParams{
UserID: userID,
})
require.NoError(t, err, "Should not error when getting user by ID")
assert.NotNil(t, result, "Result should not be nil")
assert.Len(t, result, 1, "Should return exactly one user")
assert.Equal(t, userID, result[0].ID, "Should return the correct user")
assert.NotEmpty(t, result[0].Username, "User should have a username")
})
t.Run("get user by username", func(t *testing.T) {
result, err := listOnCallUsers(ctx, ListOnCallUsersParams{
Username: username,
})
require.NoError(t, err, "Should not error when getting user by username")
assert.NotNil(t, result, "Result should not be nil")
assert.Len(t, result, 1, "Should return exactly one user")
assert.Equal(t, username, result[0].Username, "Should return the correct user")
assert.NotEmpty(t, result[0].ID, "User should have an ID")
})
t.Run("get user with invalid ID", func(t *testing.T) {
_, err := listOnCallUsers(ctx, ListOnCallUsersParams{
UserID: "invalid-user-id",
})
assert.Error(t, err, "Should error when getting user with invalid ID")
})
t.Run("get user with invalid username", func(t *testing.T) {
result, err := listOnCallUsers(ctx, ListOnCallUsersParams{
Username: "invalid-username",
})
require.NoError(t, err, "Should not error when getting user with invalid username")
assert.Empty(t, result, "Should return empty result set for invalid username")
})
}
func TestCloudGetAlertGroup(t *testing.T) {
ctx := createCloudTestContext(t, "OnCall", "GRAFANA_URL", "GRAFANA_API_KEY")
// First, get a list of alert groups to find a valid ID to test with
alertGroups, err := listAlertGroups(ctx, ListAlertGroupsParams{})
require.NoError(t, err, "Should not error when listing alert groups")
require.NotEmpty(t, alertGroups, "Should have at least one alert group to test with")
alertGroupID := alertGroups[0].ID
t.Run("get alert group by ID", func(t *testing.T) {
result, err := getAlertGroup(ctx, GetAlertGroupParams{
AlertGroupID: alertGroupID,
})
require.NoError(t, err, "Should not error when getting alert group by ID")
assert.NotNil(t, result, "Result should not be nil")
assert.Equal(t, alertGroupID, result.ID, "Should return the correct alert group")
assert.NotEmpty(t, result.Title, "Alert group should have a title")
assert.NotEmpty(t, result.State, "Alert group should have a state")
})
t.Run("get alert group with invalid ID", func(t *testing.T) {
_, err := getAlertGroup(ctx, GetAlertGroupParams{
AlertGroupID: "invalid-alert-group-id",
})
assert.Error(t, err, "Should error when getting alert group with invalid ID")
})
}
```
--------------------------------------------------------------------------------
/proxied_tools.go:
--------------------------------------------------------------------------------
```go
package mcpgrafana
import (
"context"
"fmt"
"log/slog"
"net/http"
"strings"
"sync"
"github.com/go-openapi/runtime"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// MCPDatasourceConfig defines configuration for a datasource type that supports MCP
type MCPDatasourceConfig struct {
Type string
EndpointPath string // e.g., "/api/mcp"
}
// mcpEnabledDatasources is a registry of datasource types that support MCP
var mcpEnabledDatasources = map[string]MCPDatasourceConfig{
"tempo": {Type: "tempo", EndpointPath: "/api/mcp"},
// Future: add other datasource types here
}
// DiscoveredDatasource represents a datasource that supports MCP
type DiscoveredDatasource struct {
UID string
Name string
Type string
MCPURL string // The MCP endpoint URL
}
// discoverMCPDatasources discovers datasources that support MCP
// Returns a list of datasources with MCP endpoints
func discoverMCPDatasources(ctx context.Context) ([]DiscoveredDatasource, error) {
gc := GrafanaClientFromContext(ctx)
if gc == nil {
return nil, fmt.Errorf("grafana client not found in context")
}
var discovered []DiscoveredDatasource
// List all datasources
resp, err := gc.Datasources.GetDataSources()
if err != nil {
return nil, fmt.Errorf("failed to list datasources: %w", err)
}
// Get the Grafana base URL from context
config := GrafanaConfigFromContext(ctx)
if config.URL == "" {
return nil, fmt.Errorf("grafana url not found in context")
}
grafanaBaseURL := config.URL
// Filter for datasources that support MCP
for _, ds := range resp.Payload {
// Check if this datasource type supports MCP
dsConfig, supported := mcpEnabledDatasources[ds.Type]
if !supported {
continue
}
// Check if the datasource instance has MCP enabled
// We use a DELETE request to probe the MCP endpoint since:
// - GET would start an event stream and hang
// - POST doesn't work with the Grafana OpenAPI client
// - DELETE returns 200 if MCP is enabled, 404 if not
_, err := gc.Datasources.DatasourceProxyDELETEByUIDcalls(ds.UID, strings.TrimPrefix(dsConfig.EndpointPath, "/"))
if err == nil {
// Something strange happened - the server should never return a 202 for this really. Skip.
continue
}
if apiErr, ok := err.(*runtime.APIError); !ok || (ok && !apiErr.IsCode(http.StatusOK)) {
// Not a 200 response, MCP not enabled
continue
}
// Build the MCP endpoint URL using Grafana's datasource proxy API
// Format: <grafana URL>/api/datasources/proxy/uid/<uid><endpoint_path>
mcpURL := fmt.Sprintf("%s/api/datasources/proxy/uid/%s%s", grafanaBaseURL, ds.UID, dsConfig.EndpointPath)
discovered = append(discovered, DiscoveredDatasource{
UID: ds.UID,
Name: ds.Name,
Type: ds.Type,
MCPURL: mcpURL,
})
}
slog.DebugContext(ctx, "discovered MCP datasources", "count", len(discovered))
return discovered, nil
}
// addDatasourceUidParameter adds a required datasourceUid parameter to a tool's input schema
func addDatasourceUidParameter(tool mcp.Tool, datasourceType string) mcp.Tool {
modifiedTool := tool
// Prefix tool name with datasource type (e.g., "tempo_traceql-search")
modifiedTool.Name = datasourceType + "_" + tool.Name
// Add datasourceUid to the input schema
if modifiedTool.InputSchema.Properties == nil {
modifiedTool.InputSchema.Properties = make(map[string]any)
}
modifiedTool.InputSchema.Properties["datasourceUid"] = map[string]any{
"type": "string",
"description": "UID of the " + datasourceType + " datasource to query",
}
// Add to required fields
modifiedTool.InputSchema.Required = append(modifiedTool.InputSchema.Required, "datasourceUid")
return modifiedTool
}
// parseProxiedToolName extracts datasource type and original tool name from a proxied tool name
// Format: <datasource_type>_<original_tool_name>
// Returns: datasourceType, originalToolName, error
func parseProxiedToolName(toolName string) (string, string, error) {
parts := strings.SplitN(toolName, "_", 2)
if len(parts) != 2 {
return "", "", fmt.Errorf("invalid proxied tool name format: %s", toolName)
}
return parts[0], parts[1], nil
}
// ToolManager manages proxied tools (either per-session or server-wide)
type ToolManager struct {
sm *SessionManager
server *server.MCPServer
// Whether to enable proxied tools.
enableProxiedTools bool
// For stdio transport: store clients at manager level (single-tenant).
// These will be unused for HTTP/SSE transports.
serverMode bool // true if using server-wide tools (stdio), false for per-session (HTTP/SSE)
serverClients map[string]*ProxiedClient
clientsMutex sync.RWMutex
}
// NewToolManager creates a new ToolManager
func NewToolManager(sm *SessionManager, mcpServer *server.MCPServer, opts ...toolManagerOption) *ToolManager {
tm := &ToolManager{
sm: sm,
server: mcpServer,
serverClients: make(map[string]*ProxiedClient),
}
for _, opt := range opts {
opt(tm)
}
return tm
}
type toolManagerOption func(*ToolManager)
// WithProxiedTools sets whether proxied tools are enabled
func WithProxiedTools(enabled bool) toolManagerOption {
return func(tm *ToolManager) {
tm.enableProxiedTools = enabled
}
}
// InitializeAndRegisterServerTools discovers datasources and registers tools on the server (for stdio transport)
// This should be called once at server startup for single-tenant stdio servers
func (tm *ToolManager) InitializeAndRegisterServerTools(ctx context.Context) error {
if !tm.enableProxiedTools {
return nil
}
// Mark as server mode (stdio transport)
tm.serverMode = true
// Discover datasources with MCP support
discovered, err := discoverMCPDatasources(ctx)
if err != nil {
return fmt.Errorf("failed to discover MCP datasources: %w", err)
}
if len(discovered) == 0 {
slog.Info("no MCP datasources discovered")
return nil
}
// Connect to each datasource and store in manager
tm.clientsMutex.Lock()
for _, ds := range discovered {
client, err := NewProxiedClient(ctx, ds.UID, ds.Name, ds.Type, ds.MCPURL)
if err != nil {
slog.Error("failed to create proxied client", "datasource", ds.UID, "error", err)
continue
}
key := ds.Type + "_" + ds.UID
tm.serverClients[key] = client
}
clientCount := len(tm.serverClients)
tm.clientsMutex.Unlock()
if clientCount == 0 {
slog.Warn("no proxied clients created")
return nil
}
slog.Info("connected to proxied MCP servers", "datasources", clientCount)
// Collect and register all unique tools
tm.clientsMutex.RLock()
toolMap := make(map[string]mcp.Tool)
for _, client := range tm.serverClients {
for _, tool := range client.ListTools() {
toolName := client.DatasourceType + "_" + tool.Name
if _, exists := toolMap[toolName]; !exists {
modifiedTool := addDatasourceUidParameter(tool, client.DatasourceType)
toolMap[toolName] = modifiedTool
}
}
}
tm.clientsMutex.RUnlock()
// Register tools on the server (not per-session)
for toolName, tool := range toolMap {
handler := NewProxiedToolHandler(tm.sm, tm, toolName)
tm.server.AddTool(tool, handler.Handle)
}
slog.Info("registered proxied tools on server", "tools", len(toolMap))
return nil
}
// InitializeAndRegisterProxiedTools discovers datasources, creates clients, and registers tools per-session
// This should be called in OnBeforeListTools and OnBeforeCallTool hooks for HTTP/SSE transports
func (tm *ToolManager) InitializeAndRegisterProxiedTools(ctx context.Context, session server.ClientSession) {
if !tm.enableProxiedTools {
return
}
sessionID := session.SessionID()
state, exists := tm.sm.GetSession(sessionID)
if !exists {
// Session exists in server context but not in our SessionManager yet
tm.sm.CreateSession(ctx, session)
state, exists = tm.sm.GetSession(sessionID)
if !exists {
slog.Error("failed to create session in SessionManager", "sessionID", sessionID)
return
}
}
// Step 1: Discover and connect (guaranteed to run exactly once per session)
state.initOnce.Do(func() {
// Discover datasources with MCP support
discovered, err := discoverMCPDatasources(ctx)
if err != nil {
slog.Error("failed to discover MCP datasources", "error", err)
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.mutex.Unlock()
return
}
state.mutex.Lock()
// For each discovered datasource, create a proxied client
for _, ds := range discovered {
client, err := NewProxiedClient(ctx, ds.UID, ds.Name, ds.Type, ds.MCPURL)
if err != nil {
slog.Error("failed to create proxied client", "datasource", ds.UID, "error", err)
continue
}
// Store the client
key := ds.Type + "_" + ds.UID
state.proxiedClients[key] = client
}
state.proxiedToolsInitialized = true
state.mutex.Unlock()
slog.Info("connected to proxied MCP servers", "session", sessionID, "datasources", len(state.proxiedClients))
})
// Step 2: Register tools with the MCP server
state.mutex.Lock()
defer state.mutex.Unlock()
// Check if tools already registered
if len(state.proxiedTools) > 0 {
return
}
// Check if we have any clients (discovery should have happened above)
if len(state.proxiedClients) == 0 {
return
}
// First pass: collect all unique tools and track which datasources support them
toolMap := make(map[string]mcp.Tool) // unique tools by name
for key, client := range state.proxiedClients {
remoteTools := client.ListTools()
for _, tool := range remoteTools {
// Tool name format: datasourceType_originalToolName (e.g., "tempo_traceql-search")
toolName := client.DatasourceType + "_" + tool.Name
// Store the tool if we haven't seen it yet
if _, exists := toolMap[toolName]; !exists {
// Add datasourceUid parameter to the tool
modifiedTool := addDatasourceUidParameter(tool, client.DatasourceType)
toolMap[toolName] = modifiedTool
}
// Track which datasources support this tool
state.toolToDatasources[toolName] = append(state.toolToDatasources[toolName], key)
}
}
// Second pass: register all unique tools at once (reduces listChanged notifications)
var serverTools []server.ServerTool
for toolName, tool := range toolMap {
handler := NewProxiedToolHandler(tm.sm, tm, toolName)
serverTools = append(serverTools, server.ServerTool{
Tool: tool,
Handler: handler.Handle,
})
state.proxiedTools = append(state.proxiedTools, tool)
}
if err := tm.server.AddSessionTools(sessionID, serverTools...); err != nil {
slog.Warn("failed to add session tools", "session", sessionID, "error", err)
} else {
slog.Info("registered proxied tools", "session", sessionID, "tools", len(state.proxiedTools))
}
}
// GetServerClient retrieves a proxied client from server-level storage (for stdio transport)
func (tm *ToolManager) GetServerClient(datasourceType, datasourceUID string) (*ProxiedClient, error) {
tm.clientsMutex.RLock()
defer tm.clientsMutex.RUnlock()
key := datasourceType + "_" + datasourceUID
client, exists := tm.serverClients[key]
if !exists {
// List available datasources to help with debugging
var availableUIDs []string
for _, c := range tm.serverClients {
if c.DatasourceType == datasourceType {
availableUIDs = append(availableUIDs, c.DatasourceUID)
}
}
if len(availableUIDs) > 0 {
return nil, fmt.Errorf("datasource '%s' not found. Available %s datasources: %v", datasourceUID, datasourceType, availableUIDs)
}
return nil, fmt.Errorf("datasource '%s' not found. No %s datasources with MCP support are configured", datasourceUID, datasourceType)
}
return client, nil
}
```
--------------------------------------------------------------------------------
/proxied_tools_test.go:
--------------------------------------------------------------------------------
```go
package mcpgrafana
import (
"context"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/assert"
)
func TestSessionStateRaceConditions(t *testing.T) {
t.Run("concurrent initialization with sync.Once is safe", func(t *testing.T) {
state := newSessionState()
var initCounter int32
var wg sync.WaitGroup
// Launch 100 goroutines that all try to initialize at once
const numGoroutines = 100
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func() {
defer wg.Done()
state.initOnce.Do(func() {
// Simulate initialization work
atomic.AddInt32(&initCounter, 1)
time.Sleep(10 * time.Millisecond) // Simulate some work
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.mutex.Unlock()
})
}()
}
wg.Wait()
// Verify initialization happened exactly once
assert.Equal(t, int32(1), atomic.LoadInt32(&initCounter),
"Initialization should run exactly once despite 100 concurrent calls")
assert.True(t, state.proxiedToolsInitialized)
})
t.Run("concurrent reads and writes with mutex protection", func(t *testing.T) {
state := newSessionState()
var wg sync.WaitGroup
// Writer goroutines
for i := 0; i < 10; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
state.mutex.Lock()
key := "tempo_" + string(rune('a'+id))
state.proxiedClients[key] = &ProxiedClient{
DatasourceUID: key,
DatasourceName: "Test " + key,
DatasourceType: "tempo",
}
state.mutex.Unlock()
}(i)
}
// Reader goroutines
for i := 0; i < 10; i++ {
wg.Add(1)
go func() {
defer wg.Done()
state.mutex.RLock()
_ = len(state.proxiedClients)
state.mutex.RUnlock()
}()
}
wg.Wait()
// Verify all writes succeeded
state.mutex.RLock()
count := len(state.proxiedClients)
state.mutex.RUnlock()
assert.Equal(t, 10, count, "All 10 clients should be stored")
})
t.Run("concurrent tool registration is safe", func(t *testing.T) {
state := newSessionState()
var wg sync.WaitGroup
// Multiple goroutines trying to register tools
const numGoroutines = 50
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func(id int) {
defer wg.Done()
state.mutex.Lock()
toolName := "tempo_tool-" + string(rune('a'+id%26))
if state.toolToDatasources[toolName] == nil {
state.toolToDatasources[toolName] = []string{}
}
state.toolToDatasources[toolName] = append(
state.toolToDatasources[toolName],
"datasource_"+string(rune('a'+id%26)),
)
state.mutex.Unlock()
}(i)
}
wg.Wait()
// Verify the tool mappings exist
state.mutex.RLock()
defer state.mutex.RUnlock()
assert.Greater(t, len(state.toolToDatasources), 0, "Should have tool mappings")
})
}
func TestSessionManagerConcurrency(t *testing.T) {
t.Run("concurrent session creation is safe", func(t *testing.T) {
sm := NewSessionManager()
var wg sync.WaitGroup
// Create many sessions concurrently
const numSessions = 100
wg.Add(numSessions)
for i := 0; i < numSessions; i++ {
go func(id int) {
defer wg.Done()
sessionID := "session-" + string(rune('a'+id%26)) + "-" + string(rune('0'+id/26))
mockSession := &mockClientSession{id: sessionID}
sm.CreateSession(context.Background(), mockSession)
}(i)
}
wg.Wait()
// Verify all sessions were created
sm.mutex.RLock()
count := len(sm.sessions)
sm.mutex.RUnlock()
assert.Equal(t, numSessions, count, "All sessions should be created")
})
t.Run("concurrent get and remove is safe", func(t *testing.T) {
sm := NewSessionManager()
// Pre-populate sessions
for i := 0; i < 50; i++ {
sessionID := "session-" + string(rune('a'+i%26))
mockSession := &mockClientSession{id: sessionID}
sm.CreateSession(context.Background(), mockSession)
}
var wg sync.WaitGroup
// Readers
for i := 0; i < 50; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
sessionID := "session-" + string(rune('a'+id%26))
_, _ = sm.GetSession(sessionID)
}(i)
}
// Writers (removers)
for i := 0; i < 25; i++ {
wg.Add(1)
go func(id int) {
defer wg.Done()
sessionID := "session-" + string(rune('a'+id%26))
mockSession := &mockClientSession{id: sessionID}
sm.RemoveSession(context.Background(), mockSession)
}(i)
}
wg.Wait()
// Test passed if no race conditions occurred
})
}
func TestInitOncePattern(t *testing.T) {
t.Run("verify sync.Once guarantees single execution", func(t *testing.T) {
var once sync.Once
var counter int32
var wg sync.WaitGroup
// Simulate what happens in InitializeAndRegisterProxiedTools
initFunc := func() {
atomic.AddInt32(&counter, 1)
// Simulate expensive initialization
time.Sleep(50 * time.Millisecond)
}
// Launch many concurrent calls
for i := 0; i < 1000; i++ {
wg.Add(1)
go func() {
defer wg.Done()
once.Do(initFunc)
}()
}
wg.Wait()
assert.Equal(t, int32(1), atomic.LoadInt32(&counter),
"sync.Once should guarantee function runs exactly once")
})
t.Run("sync.Once with different functions only runs first", func(t *testing.T) {
var once sync.Once
var result string
var mu sync.Mutex
once.Do(func() {
mu.Lock()
result = "first"
mu.Unlock()
})
once.Do(func() {
mu.Lock()
result = "second"
mu.Unlock()
})
mu.Lock()
finalResult := result
mu.Unlock()
assert.Equal(t, "first", finalResult, "Only first function should execute")
})
}
func TestProxiedToolsInitializationFlow(t *testing.T) {
t.Run("initialization state transitions are correct", func(t *testing.T) {
state := newSessionState()
// Initial state
assert.False(t, state.proxiedToolsInitialized)
assert.Empty(t, state.proxiedClients)
assert.Empty(t, state.proxiedTools)
// Simulate initialization
state.initOnce.Do(func() {
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.proxiedClients["tempo_test"] = &ProxiedClient{
DatasourceUID: "test",
DatasourceName: "Test",
DatasourceType: "tempo",
}
state.mutex.Unlock()
})
// Verify state after initialization
state.mutex.RLock()
initialized := state.proxiedToolsInitialized
clientCount := len(state.proxiedClients)
state.mutex.RUnlock()
assert.True(t, initialized)
assert.Equal(t, 1, clientCount)
})
t.Run("multiple sessions maintain separate state", func(t *testing.T) {
sm := NewSessionManager()
// Create two sessions
session1 := &mockClientSession{id: "session-1"}
session2 := &mockClientSession{id: "session-2"}
sm.CreateSession(context.Background(), session1)
sm.CreateSession(context.Background(), session2)
state1, _ := sm.GetSession("session-1")
state2, _ := sm.GetSession("session-2")
// Initialize only session1
state1.initOnce.Do(func() {
state1.mutex.Lock()
state1.proxiedToolsInitialized = true
state1.mutex.Unlock()
})
// Verify states are independent
assert.True(t, state1.proxiedToolsInitialized)
assert.False(t, state2.proxiedToolsInitialized)
assert.NotSame(t, state1, state2)
})
}
func TestRaceConditionDemonstration(t *testing.T) {
t.Run("old pattern WITHOUT sync.Once would have race condition", func(t *testing.T) {
// This test demonstrates what WOULD happen with the old mutex-check pattern
state := newSessionState()
var discoveryCallCount int32
var wg sync.WaitGroup
// Simulate the OLD pattern (mutex check, unlock, then do work)
oldPatternInitialize := func() {
state.mutex.Lock()
// Check if already initialized
if state.proxiedToolsInitialized {
state.mutex.Unlock()
return
}
alreadyDiscovered := state.proxiedToolsInitialized
state.mutex.Unlock() // ❌ OLD PATTERN: Unlock before expensive work
if !alreadyDiscovered {
// Simulate discovery work that should only happen once
atomic.AddInt32(&discoveryCallCount, 1)
time.Sleep(10 * time.Millisecond) // Simulate expensive operation
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.mutex.Unlock()
}
}
// Launch concurrent initializations
const numGoroutines = 10
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func() {
defer wg.Done()
oldPatternInitialize()
}()
}
wg.Wait()
// With the old pattern, multiple goroutines can get past the check
// and call discovery multiple times
count := atomic.LoadInt32(&discoveryCallCount)
if count > 1 {
t.Logf("OLD PATTERN: Discovery called %d times (race condition!)", count)
}
// We can't assert > 1 reliably because timing matters, but this demonstrates the problem
})
t.Run("new pattern WITH sync.Once prevents race condition", func(t *testing.T) {
// This test demonstrates the NEW pattern with sync.Once
state := newSessionState()
var discoveryCallCount int32
var wg sync.WaitGroup
// NEW pattern: sync.Once guarantees single execution
newPatternInitialize := func() {
state.initOnce.Do(func() {
// Simulate discovery work that should only happen once
atomic.AddInt32(&discoveryCallCount, 1)
time.Sleep(10 * time.Millisecond) // Simulate expensive operation
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.mutex.Unlock()
})
}
// Launch concurrent initializations
const numGoroutines = 10
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func() {
defer wg.Done()
newPatternInitialize()
}()
}
wg.Wait()
// With sync.Once, discovery is guaranteed to run exactly once
count := atomic.LoadInt32(&discoveryCallCount)
assert.Equal(t, int32(1), count, "NEW PATTERN: Discovery must be called exactly once")
})
}
func TestRaceDetector(t *testing.T) {
// This test is primarily valuable when run with -race flag
t.Run("stress test with race detector", func(t *testing.T) {
sm := NewSessionManager()
var wg sync.WaitGroup
// Create a mix of operations happening concurrently
for i := 0; i < 20; i++ {
sessionID := "stress-session-" + string(rune('a'+i%10))
// Create session
wg.Add(1)
go func(sid string) {
defer wg.Done()
mockSession := &mockClientSession{id: sid}
sm.CreateSession(context.Background(), mockSession)
}(sessionID)
// Initialize session state
wg.Add(1)
go func(sid string) {
defer wg.Done()
time.Sleep(time.Millisecond) // Let creation happen first
state, exists := sm.GetSession(sid)
if exists {
state.initOnce.Do(func() {
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.mutex.Unlock()
})
}
}(sessionID)
// Read session state
wg.Add(1)
go func(sid string) {
defer wg.Done()
time.Sleep(2 * time.Millisecond)
state, exists := sm.GetSession(sid)
if exists {
state.mutex.RLock()
_ = state.proxiedToolsInitialized
state.mutex.RUnlock()
}
}(sessionID)
}
wg.Wait()
// If we get here without race detector warnings, we're good
t.Log("Stress test completed without race conditions")
})
}
// mockClientSession implements server.ClientSession for testing
type mockClientSession struct {
id string
notifChannel chan mcp.JSONRPCNotification
isInitialized bool
}
func (m *mockClientSession) SessionID() string {
return m.id
}
func (m *mockClientSession) NotificationChannel() chan<- mcp.JSONRPCNotification {
if m.notifChannel == nil {
m.notifChannel = make(chan mcp.JSONRPCNotification, 10)
}
return m.notifChannel
}
func (m *mockClientSession) Initialize() {
m.isInitialized = true
}
func (m *mockClientSession) Initialized() bool {
return m.isInitialized
}
```
--------------------------------------------------------------------------------
/tools/dashboard_test.go:
--------------------------------------------------------------------------------
```go
// Requires a Grafana instance running on localhost:3000,
// with a dashboard provisioned.
// Run with `go test -tags integration`.
//go:build integration
package tools
import (
"context"
"testing"
"github.com/grafana/grafana-openapi-client-go/models"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const (
newTestDashboardName = "Integration Test"
)
// getExistingDashboardUID will fetch an existing dashboard for test purposes
// It will search for exisiting dashboards and return the first, otherwise
// will trigger a test error
func getExistingTestDashboard(t *testing.T, ctx context.Context, dashboardName string) *models.Hit {
// Make sure we query for the existing dashboard, not a folder
if dashboardName == "" {
dashboardName = "Demo"
}
searchResults, err := searchDashboards(ctx, SearchDashboardsParams{
Query: dashboardName,
})
require.NoError(t, err)
require.Greater(t, len(searchResults), 0, "No dashboards found")
return searchResults[0]
}
// getExistingTestDashboardJSON will fetch the JSON map for an existing
// dashboard in the test environment
func getTestDashboardJSON(t *testing.T, ctx context.Context, dashboard *models.Hit) map[string]interface{} {
result, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := result.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
return dashboardMap
}
func TestDashboardTools(t *testing.T) {
t.Run("get dashboard by uid", func(t *testing.T) {
ctx := newTestContext()
// First, let's search for a dashboard to get its UID
dashboard := getExistingTestDashboard(t, ctx, "")
// Now test the get dashboard by uid functionality
result, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := result.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
assert.Equal(t, dashboard.UID, dashboardMap["uid"])
assert.NotNil(t, result.Meta)
})
t.Run("get dashboard by uid - invalid uid", func(t *testing.T) {
ctx := newTestContext()
_, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: "non-existent-uid",
})
require.Error(t, err)
})
t.Run("update dashboard - create new", func(t *testing.T) {
ctx := newTestContext()
// Get the dashboard JSON
// In this case, we will create a new dashboard with the same
// content but different Title, and disable "overwrite"
dashboard := getExistingTestDashboard(t, ctx, "")
dashboardMap := getTestDashboardJSON(t, ctx, dashboard)
// Avoid a clash by unsetting the existing IDs
delete(dashboardMap, "uid")
delete(dashboardMap, "id")
// Set a new title and tag
dashboardMap["title"] = newTestDashboardName
dashboardMap["tags"] = []string{"integration-test"}
params := UpdateDashboardParams{
Dashboard: dashboardMap,
Message: "creating a new dashboard",
Overwrite: false,
UserID: 1,
}
// Only pass in the Folder UID if it exists
if dashboard.FolderUID != "" {
params.FolderUID = dashboard.FolderUID
}
// create the dashboard
_, err := updateDashboard(ctx, params)
require.NoError(t, err)
})
t.Run("update dashboard - overwrite existing", func(t *testing.T) {
ctx := newTestContext()
// Get the dashboard JSON for the non-provisioned dashboard we've created
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
dashboardMap := getTestDashboardJSON(t, ctx, dashboard)
params := UpdateDashboardParams{
Dashboard: dashboardMap,
Message: "updating existing dashboard",
Overwrite: true,
UserID: 1,
}
// Only pass in the Folder UID if it exists
if dashboard.FolderUID != "" {
params.FolderUID = dashboard.FolderUID
}
// update the dashboard
_, err := updateDashboard(ctx, params)
require.NoError(t, err)
})
t.Run("get dashboard panel queries", func(t *testing.T) {
ctx := newTestContext()
// Get the test dashboard
dashboard := getExistingTestDashboard(t, ctx, "")
result, err := GetDashboardPanelQueriesTool(ctx, DashboardPanelQueriesParams{
UID: dashboard.UID,
})
require.NoError(t, err)
assert.Greater(t, len(result), 0, "Should return at least one panel query")
// The initial demo dashboard plus for all dashboards created by the integration tests,
// every panel should have identical title and query values.
// Datasource UID may differ. Datasource type can be an empty string as well but on the demo and test dashboards it should be "prometheus".
for _, panelQuery := range result {
assert.Equal(t, panelQuery.Title, "Node Load")
assert.Equal(t, panelQuery.Query, "node_load1")
assert.NotEmpty(t, panelQuery.Datasource.UID)
assert.Equal(t, panelQuery.Datasource.Type, "prometheus")
}
})
// Tests for new Issue #101 context window management tools
t.Run("get dashboard summary", func(t *testing.T) {
ctx := newTestContext()
// Get the test dashboard
dashboard := getExistingTestDashboard(t, ctx, "")
result, err := getDashboardSummary(ctx, GetDashboardSummaryParams{
UID: dashboard.UID,
})
require.NoError(t, err)
assert.Equal(t, dashboard.UID, result.UID)
assert.NotEmpty(t, result.Title)
assert.Greater(t, result.PanelCount, 0, "Should have at least one panel")
assert.Len(t, result.Panels, result.PanelCount, "Panel count should match panels array length")
assert.NotNil(t, result.Meta)
// Check that panels have expected structure
for _, panel := range result.Panels {
assert.NotEmpty(t, panel.Title)
assert.NotEmpty(t, panel.Type)
assert.GreaterOrEqual(t, panel.QueryCount, 0)
}
})
t.Run("get dashboard property - title", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, "")
result, err := getDashboardProperty(ctx, GetDashboardPropertyParams{
UID: dashboard.UID,
JSONPath: "$.title",
})
require.NoError(t, err)
title, ok := result.(string)
require.True(t, ok, "Title should be a string")
assert.NotEmpty(t, title)
})
t.Run("get dashboard property - panel titles", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, "")
result, err := getDashboardProperty(ctx, GetDashboardPropertyParams{
UID: dashboard.UID,
JSONPath: "$.panels[*].title",
})
require.NoError(t, err)
titles, ok := result.([]interface{})
require.True(t, ok, "Panel titles should be an array")
assert.Greater(t, len(titles), 0, "Should have at least one panel title")
for _, title := range titles {
titleStr, ok := title.(string)
require.True(t, ok, "Each title should be a string")
assert.NotEmpty(t, titleStr)
}
})
t.Run("get dashboard property - invalid path", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, "")
_, err := getDashboardProperty(ctx, GetDashboardPropertyParams{
UID: dashboard.UID,
JSONPath: "$.nonexistent.path",
})
require.Error(t, err, "Should fail for non-existent path")
})
t.Run("update dashboard - patch title", func(t *testing.T) {
ctx := newTestContext()
// Get our test dashboard (not the provisioned one)
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
newTitle := "Updated Integration Test Dashboard"
result, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "replace",
Path: "$.title",
Value: newTitle,
},
},
Message: "Updated title via patch",
})
require.NoError(t, err)
assert.NotNil(t, result)
// Verify the change was applied
updatedDashboard, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := updatedDashboard.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
assert.Equal(t, newTitle, dashboardMap["title"])
})
t.Run("update dashboard - patch add description", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
description := "This is a test description added via patch"
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "add",
Path: "$.description",
Value: description,
},
},
Message: "Added description via patch",
})
require.NoError(t, err)
// Verify the description was added
updatedDashboard, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := updatedDashboard.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
assert.Equal(t, description, dashboardMap["description"])
})
t.Run("update dashboard - patch remove description", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "remove",
Path: "$.description",
},
},
Message: "Removed description via patch",
})
require.NoError(t, err)
// Verify the description was removed
updatedDashboard, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := updatedDashboard.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
_, hasDescription := dashboardMap["description"]
assert.False(t, hasDescription, "Description should be removed")
})
t.Run("update dashboard - unsupported operation", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "copy", // Unsupported operation
Path: "$.title",
Value: "New Title",
},
},
})
require.Error(t, err, "Should fail for unsupported operation")
})
t.Run("update dashboard - invalid parameters", func(t *testing.T) {
ctx := newTestContext()
_, err := updateDashboard(ctx, UpdateDashboardParams{
// Neither dashboard nor (uid + operations) provided
})
require.Error(t, err, "Should fail when no valid parameters provided")
})
t.Run("update dashboard - append to panels array", func(t *testing.T) {
ctx := newTestContext()
// Get our test dashboard
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
// Create a new panel to append
newPanel := map[string]interface{}{
"id": 999,
"title": "New Appended Panel",
"type": "stat",
"targets": []interface{}{
map[string]interface{}{
"expr": "up",
},
},
"gridPos": map[string]interface{}{
"h": 8,
"w": 12,
"x": 0,
"y": 8,
},
}
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "add",
Path: "$.panels/-",
Value: newPanel,
},
},
Message: "Appended new panel via /- syntax",
})
require.NoError(t, err)
// Verify the panel was appended
updatedDashboard, err := getDashboardByUID(ctx, GetDashboardByUIDParams{
UID: dashboard.UID,
})
require.NoError(t, err)
dashboardMap, ok := updatedDashboard.Dashboard.(map[string]interface{})
require.True(t, ok, "Dashboard should be a map")
panels, ok := dashboardMap["panels"].([]interface{})
require.True(t, ok, "Panels should be an array")
// Check that the new panel was appended (should be the last panel)
lastPanel, ok := panels[len(panels)-1].(map[string]interface{})
require.True(t, ok, "Last panel should be an object")
assert.Equal(t, "New Appended Panel", lastPanel["title"])
assert.Equal(t, float64(999), lastPanel["id"]) // JSON unmarshaling converts to float64
})
t.Run("update dashboard - remove with append syntax should fail", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "remove",
Path: "$.panels/-", // Invalid: remove with append syntax
},
},
})
require.Error(t, err, "Should fail when using remove operation with append syntax")
})
t.Run("update dashboard - append to non-array should fail", func(t *testing.T) {
ctx := newTestContext()
dashboard := getExistingTestDashboard(t, ctx, newTestDashboardName)
_, err := updateDashboard(ctx, UpdateDashboardParams{
UID: dashboard.UID,
Operations: []PatchOperation{
{
Op: "add",
Path: "$.title/-", // Invalid: title is not an array
Value: "Invalid",
},
},
})
require.Error(t, err, "Should fail when trying to append to non-array field")
})
}
```
--------------------------------------------------------------------------------
/tools_test.go:
--------------------------------------------------------------------------------
```go
//go:build unit
// +build unit
package mcpgrafana
import (
"context"
"errors"
"testing"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
type testToolParams struct {
Name string `json:"name" jsonschema:"required,description=The name parameter"`
Value int `json:"value" jsonschema:"required,description=The value parameter"`
Optional bool `json:"optional,omitempty" jsonschema:"description=An optional parameter"`
}
func testToolHandler(ctx context.Context, params testToolParams) (*mcp.CallToolResult, error) {
if params.Name == "error" {
return nil, errors.New("test error")
}
return mcp.NewToolResultText(params.Name + ": " + string(rune(params.Value))), nil
}
type emptyToolParams struct{}
func emptyToolHandler(ctx context.Context, params emptyToolParams) (*mcp.CallToolResult, error) {
return mcp.NewToolResultText("empty"), nil
}
// New handlers for different return types
func stringToolHandler(ctx context.Context, params testToolParams) (string, error) {
if params.Name == "error" {
return "", errors.New("test error")
}
if params.Name == "empty" {
return "", nil
}
return params.Name + ": " + string(rune(params.Value)), nil
}
func stringPtrToolHandler(ctx context.Context, params testToolParams) (*string, error) {
if params.Name == "error" {
return nil, errors.New("test error")
}
if params.Name == "nil" {
return nil, nil
}
if params.Name == "empty" {
empty := ""
return &empty, nil
}
result := params.Name + ": " + string(rune(params.Value))
return &result, nil
}
type TestResult struct {
Name string `json:"name"`
Value int `json:"value"`
}
func structToolHandler(ctx context.Context, params testToolParams) (TestResult, error) {
if params.Name == "error" {
return TestResult{}, errors.New("test error")
}
return TestResult{
Name: params.Name,
Value: params.Value,
}, nil
}
func structPtrToolHandler(ctx context.Context, params testToolParams) (*TestResult, error) {
if params.Name == "error" {
return nil, errors.New("test error")
}
if params.Name == "nil" {
return nil, nil
}
return &TestResult{
Name: params.Name,
Value: params.Value,
}, nil
}
func TestConvertTool(t *testing.T) {
t.Run("valid handler conversion", func(t *testing.T) {
tool, handler, err := ConvertTool("test_tool", "A test tool", testToolHandler)
require.NoError(t, err)
require.NotNil(t, tool)
require.NotNil(t, handler)
// Check tool properties
assert.Equal(t, "test_tool", tool.Name)
assert.Equal(t, "A test tool", tool.Description)
// Check schema properties
assert.Equal(t, "object", tool.InputSchema.Type)
assert.Contains(t, tool.InputSchema.Properties, "name")
assert.Contains(t, tool.InputSchema.Properties, "value")
assert.Contains(t, tool.InputSchema.Properties, "optional")
// Test handler execution
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "test_tool",
Arguments: map[string]any{
"name": "test",
"value": 65, // ASCII 'A'
},
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Equal(t, "test: A", resultString.Text)
// Test error handling
errorRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "test_tool",
Arguments: map[string]any{
"name": "error",
"value": 66,
},
},
}
_, err = handler(ctx, errorRequest)
assert.Error(t, err)
assert.Equal(t, "test error", err.Error())
})
t.Run("empty handler params", func(t *testing.T) {
tool, handler, err := ConvertTool("empty", "description", emptyToolHandler)
require.NoError(t, err)
require.NotNil(t, tool)
require.NotNil(t, handler)
// Check tool properties
assert.Equal(t, "empty", tool.Name)
assert.Equal(t, "description", tool.Description)
// Check schema properties
assert.Equal(t, "object", tool.InputSchema.Type)
assert.Len(t, tool.InputSchema.Properties, 0)
// Test handler execution
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "empty",
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Equal(t, "empty", resultString.Text)
})
t.Run("string return type", func(t *testing.T) {
_, handler, err := ConvertTool("string_tool", "A string tool", stringToolHandler)
require.NoError(t, err)
// Test normal string return
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_tool",
Arguments: map[string]any{
"name": "test",
"value": 65, // ASCII 'A'
},
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Equal(t, "test: A", resultString.Text)
// Test empty string return
emptyRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_tool",
Arguments: map[string]any{
"name": "empty",
"value": 65,
},
},
}
result, err = handler(ctx, emptyRequest)
require.NoError(t, err)
assert.Nil(t, result)
// Test error return
errorRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_tool",
Arguments: map[string]any{
"name": "error",
"value": 65,
},
},
}
_, err = handler(ctx, errorRequest)
assert.Error(t, err)
assert.Equal(t, "test error", err.Error())
})
t.Run("string pointer return type", func(t *testing.T) {
_, handler, err := ConvertTool("string_ptr_tool", "A string pointer tool", stringPtrToolHandler)
require.NoError(t, err)
// Test normal string pointer return
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_ptr_tool",
Arguments: map[string]any{
"name": "test",
"value": 65, // ASCII 'A'
},
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Equal(t, "test: A", resultString.Text)
// Test nil string pointer return
nilRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_ptr_tool",
Arguments: map[string]any{
"name": "nil",
"value": 65,
},
},
}
result, err = handler(ctx, nilRequest)
require.NoError(t, err)
assert.Nil(t, result)
// Test empty string pointer return
emptyRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_ptr_tool",
Arguments: map[string]any{
"name": "empty",
"value": 65,
},
},
}
result, err = handler(ctx, emptyRequest)
require.NoError(t, err)
assert.Nil(t, result)
// Test error return
errorRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "string_ptr_tool",
Arguments: map[string]any{
"name": "error",
"value": 65,
},
},
}
_, err = handler(ctx, errorRequest)
assert.Error(t, err)
assert.Equal(t, "test error", err.Error())
})
t.Run("struct return type", func(t *testing.T) {
_, handler, err := ConvertTool("struct_tool", "A struct tool", structToolHandler)
require.NoError(t, err)
// Test normal struct return
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "struct_tool",
Arguments: map[string]any{
"name": "test",
"value": 65, // ASCII 'A'
},
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Contains(t, resultString.Text, `"name":"test"`)
assert.Contains(t, resultString.Text, `"value":65`)
// Test error return
errorRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "struct_tool",
Arguments: map[string]any{
"name": "error",
"value": 65,
},
},
}
_, err = handler(ctx, errorRequest)
assert.Error(t, err)
assert.Equal(t, "test error", err.Error())
})
t.Run("struct pointer return type", func(t *testing.T) {
_, handler, err := ConvertTool("struct_ptr_tool", "A struct pointer tool", structPtrToolHandler)
require.NoError(t, err)
// Test normal struct pointer return
ctx := context.Background()
request := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "struct_ptr_tool",
Arguments: map[string]any{
"name": "test",
"value": 65, // ASCII 'A'
},
},
}
result, err := handler(ctx, request)
require.NoError(t, err)
require.NotNil(t, result)
require.Len(t, result.Content, 1)
resultString, ok := result.Content[0].(mcp.TextContent)
require.True(t, ok)
assert.Contains(t, resultString.Text, `"name":"test"`)
assert.Contains(t, resultString.Text, `"value":65`)
// Test nil struct pointer return
nilRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "struct_ptr_tool",
Arguments: map[string]any{
"name": "nil",
"value": 65,
},
},
}
result, err = handler(ctx, nilRequest)
require.NoError(t, err)
assert.Nil(t, result)
// Test error return
errorRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Name: "struct_ptr_tool",
Arguments: map[string]any{
"name": "error",
"value": 65,
},
},
}
_, err = handler(ctx, errorRequest)
assert.Error(t, err)
assert.Equal(t, "test error", err.Error())
})
t.Run("invalid handler types", func(t *testing.T) {
// Test wrong second argument type (not a struct)
wrongSecondArgFunc := func(ctx context.Context, s string) (*mcp.CallToolResult, error) {
return nil, nil
}
_, _, err := ConvertTool("invalid", "description", wrongSecondArgFunc)
assert.Error(t, err)
assert.Contains(t, err.Error(), "second argument must be a struct")
})
t.Run("handler execution with invalid arguments", func(t *testing.T) {
_, handler, err := ConvertTool("test_tool", "A test tool", testToolHandler)
require.NoError(t, err)
// Test with invalid JSON
invalidRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Arguments: map[string]any{
"name": make(chan int), // Channels can't be marshaled to JSON
},
},
}
_, err = handler(context.Background(), invalidRequest)
assert.Error(t, err)
assert.Contains(t, err.Error(), "marshal args")
// Test with type mismatch
mismatchRequest := mcp.CallToolRequest{
Params: struct {
Name string `json:"name"`
Arguments any `json:"arguments,omitempty"`
Meta *mcp.Meta `json:"_meta,omitempty"`
}{
Arguments: map[string]any{
"name": 123, // Should be a string
"value": "not an int",
},
},
}
_, err = handler(context.Background(), mismatchRequest)
assert.Error(t, err)
assert.Contains(t, err.Error(), "unmarshal args")
})
}
func TestCreateJSONSchemaFromHandler(t *testing.T) {
schema := createJSONSchemaFromHandler(testToolHandler)
assert.Equal(t, "object", schema.Type)
assert.Len(t, schema.Required, 2) // name and value are required, optional is not
// Check properties
nameProperty, ok := schema.Properties.Get("name")
assert.True(t, ok)
assert.Equal(t, "string", nameProperty.Type)
assert.Equal(t, "The name parameter", nameProperty.Description)
valueProperty, ok := schema.Properties.Get("value")
assert.True(t, ok)
assert.Equal(t, "integer", valueProperty.Type)
assert.Equal(t, "The value parameter", valueProperty.Description)
optionalProperty, ok := schema.Properties.Get("optional")
assert.True(t, ok)
assert.Equal(t, "boolean", optionalProperty.Type)
assert.Equal(t, "An optional parameter", optionalProperty.Description)
}
```
--------------------------------------------------------------------------------
/cmd/mcp-grafana/main.go:
--------------------------------------------------------------------------------
```go
package main
import (
"context"
"errors"
"flag"
"fmt"
"log/slog"
"net/http"
"os"
"os/signal"
"slices"
"strings"
"syscall"
"time"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
mcpgrafana "github.com/grafana/mcp-grafana"
"github.com/grafana/mcp-grafana/tools"
)
func maybeAddTools(s *server.MCPServer, tf func(*server.MCPServer), enabledTools []string, disable bool, category string) {
if !slices.Contains(enabledTools, category) {
slog.Debug("Not enabling tools", "category", category)
return
}
if disable {
slog.Info("Disabling tools", "category", category)
return
}
slog.Debug("Enabling tools", "category", category)
tf(s)
}
// disabledTools indicates whether each category of tools should be disabled.
type disabledTools struct {
enabledTools string
search, datasource, incident,
prometheus, loki, alerting,
dashboard, folder, oncall, asserts, sift, admin,
pyroscope, navigation, proxied, annotations, write bool
}
// Configuration for the Grafana client.
type grafanaConfig struct {
// Whether to enable debug mode for the Grafana transport.
debug bool
// TLS configuration
tlsCertFile string
tlsKeyFile string
tlsCAFile string
tlsSkipVerify bool
}
func (dt *disabledTools) addFlags() {
flag.StringVar(&dt.enabledTools, "enabled-tools", "search,datasource,incident,prometheus,loki,alerting,dashboard,folder,oncall,asserts,sift,admin,pyroscope,navigation,proxied,annotations", "A comma separated list of tools enabled for this server. Can be overwritten entirely or by disabling specific components, e.g. --disable-search.")
flag.BoolVar(&dt.search, "disable-search", false, "Disable search tools")
flag.BoolVar(&dt.datasource, "disable-datasource", false, "Disable datasource tools")
flag.BoolVar(&dt.incident, "disable-incident", false, "Disable incident tools")
flag.BoolVar(&dt.prometheus, "disable-prometheus", false, "Disable prometheus tools")
flag.BoolVar(&dt.loki, "disable-loki", false, "Disable loki tools")
flag.BoolVar(&dt.alerting, "disable-alerting", false, "Disable alerting tools")
flag.BoolVar(&dt.dashboard, "disable-dashboard", false, "Disable dashboard tools")
flag.BoolVar(&dt.folder, "disable-folder", false, "Disable folder tools")
flag.BoolVar(&dt.oncall, "disable-oncall", false, "Disable oncall tools")
flag.BoolVar(&dt.asserts, "disable-asserts", false, "Disable asserts tools")
flag.BoolVar(&dt.sift, "disable-sift", false, "Disable sift tools")
flag.BoolVar(&dt.admin, "disable-admin", false, "Disable admin tools")
flag.BoolVar(&dt.pyroscope, "disable-pyroscope", false, "Disable pyroscope tools")
flag.BoolVar(&dt.navigation, "disable-navigation", false, "Disable navigation tools")
flag.BoolVar(&dt.proxied, "disable-proxied", false, "Disable proxied tools (tools from external MCP servers)")
flag.BoolVar(&dt.write, "disable-write", false, "Disable write tools (create/update operations)")
flag.BoolVar(&dt.annotations, "disable-annotations", false, "Disable annotation tools")
}
func (gc *grafanaConfig) addFlags() {
flag.BoolVar(&gc.debug, "debug", false, "Enable debug mode for the Grafana transport")
// TLS configuration flags
flag.StringVar(&gc.tlsCertFile, "tls-cert-file", "", "Path to TLS certificate file for client authentication")
flag.StringVar(&gc.tlsKeyFile, "tls-key-file", "", "Path to TLS private key file for client authentication")
flag.StringVar(&gc.tlsCAFile, "tls-ca-file", "", "Path to TLS CA certificate file for server verification")
flag.BoolVar(&gc.tlsSkipVerify, "tls-skip-verify", false, "Skip TLS certificate verification (insecure)")
}
func (dt *disabledTools) addTools(s *server.MCPServer) {
enabledTools := strings.Split(dt.enabledTools, ",")
enableWriteTools := !dt.write
maybeAddTools(s, tools.AddSearchTools, enabledTools, dt.search, "search")
maybeAddTools(s, tools.AddDatasourceTools, enabledTools, dt.datasource, "datasource")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddIncidentTools(mcp, enableWriteTools) }, enabledTools, dt.incident, "incident")
maybeAddTools(s, tools.AddPrometheusTools, enabledTools, dt.prometheus, "prometheus")
maybeAddTools(s, tools.AddLokiTools, enabledTools, dt.loki, "loki")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddAlertingTools(mcp, enableWriteTools) }, enabledTools, dt.alerting, "alerting")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddDashboardTools(mcp, enableWriteTools) }, enabledTools, dt.dashboard, "dashboard")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddFolderTools(mcp, enableWriteTools) }, enabledTools, dt.folder, "folder")
maybeAddTools(s, tools.AddOnCallTools, enabledTools, dt.oncall, "oncall")
maybeAddTools(s, tools.AddAssertsTools, enabledTools, dt.asserts, "asserts")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddSiftTools(mcp, enableWriteTools) }, enabledTools, dt.sift, "sift")
maybeAddTools(s, tools.AddAdminTools, enabledTools, dt.admin, "admin")
maybeAddTools(s, tools.AddPyroscopeTools, enabledTools, dt.pyroscope, "pyroscope")
maybeAddTools(s, tools.AddNavigationTools, enabledTools, dt.navigation, "navigation")
maybeAddTools(s, func(mcp *server.MCPServer) { tools.AddAnnotationTools(mcp, enableWriteTools) }, enabledTools, dt.annotations, "annotations")
}
func newServer(transport string, dt disabledTools) (*server.MCPServer, *mcpgrafana.ToolManager) {
sm := mcpgrafana.NewSessionManager()
// Declare variable for ToolManager that will be initialized after server creation
var stm *mcpgrafana.ToolManager
// Create hooks
hooks := &server.Hooks{
OnRegisterSession: []server.OnRegisterSessionHookFunc{sm.CreateSession},
OnUnregisterSession: []server.OnUnregisterSessionHookFunc{sm.RemoveSession},
}
// Add proxied tools hooks if enabled and we're not running in stdio mode.
// (stdio mode is handled by InitializeAndRegisterServerTools; per-session tools
// are not supported).
if transport != "stdio" && !dt.proxied {
// OnBeforeListTools: Discover, connect, and register tools
hooks.OnBeforeListTools = []server.OnBeforeListToolsFunc{
func(ctx context.Context, id any, request *mcp.ListToolsRequest) {
if stm != nil {
if session := server.ClientSessionFromContext(ctx); session != nil {
stm.InitializeAndRegisterProxiedTools(ctx, session)
}
}
},
}
// OnBeforeCallTool: Fallback in case client calls tool without listing first
hooks.OnBeforeCallTool = []server.OnBeforeCallToolFunc{
func(ctx context.Context, id any, request *mcp.CallToolRequest) {
if stm != nil {
if session := server.ClientSessionFromContext(ctx); session != nil {
stm.InitializeAndRegisterProxiedTools(ctx, session)
}
}
},
}
}
s := server.NewMCPServer("mcp-grafana", mcpgrafana.Version(),
server.WithInstructions(`
This server provides access to your Grafana instance and the surrounding ecosystem.
Available Capabilities:
- Dashboards: Search, retrieve, update, and create dashboards. Extract panel queries and datasource information.
- Datasources: List and fetch details for datasources.
- Prometheus & Loki: Run PromQL and LogQL queries, retrieve metric/log metadata, and explore label names/values.
- Incidents: Search, create, update, and resolve incidents in Grafana Incident.
- Sift Investigations: Start and manage Sift investigations, analyze logs/traces, find error patterns, and detect slow requests.
- Alerting: List and fetch alert rules and notification contact points.
- OnCall: View and manage on-call schedules, shifts, teams, and users.
- Admin: List teams and perform administrative tasks.
- Pyroscope: Profile applications and fetch profiling data.
- Navigation: Generate deeplink URLs for Grafana resources like dashboards, panels, and Explore queries.
- Proxied Tools: Access tools from external MCP servers (like Tempo) through dynamic discovery.
Note that some of these capabilities may be disabled. Do not try to use features that are not available via tools.
`),
server.WithHooks(hooks),
)
// Initialize ToolManager now that server is created
stm = mcpgrafana.NewToolManager(sm, s, mcpgrafana.WithProxiedTools(!dt.proxied))
dt.addTools(s)
return s, stm
}
type tlsConfig struct {
certFile, keyFile string
}
func (tc *tlsConfig) addFlags() {
flag.StringVar(&tc.certFile, "server.tls-cert-file", "", "Path to TLS certificate file for server HTTPS (required for TLS)")
flag.StringVar(&tc.keyFile, "server.tls-key-file", "", "Path to TLS private key file for server HTTPS (required for TLS)")
}
// httpServer represents a server with Start and Shutdown methods
type httpServer interface {
Start(addr string) error
Shutdown(ctx context.Context) error
}
// runHTTPServer handles the common logic for running HTTP-based servers
func runHTTPServer(ctx context.Context, srv httpServer, addr, transportName string) error {
// Start server in a goroutine
serverErr := make(chan error, 1)
go func() {
if err := srv.Start(addr); err != nil {
serverErr <- err
}
close(serverErr)
}()
// Wait for either server error or shutdown signal
select {
case err := <-serverErr:
return err
case <-ctx.Done():
slog.Info(fmt.Sprintf("%s server shutting down...", transportName))
// Create a timeout context for shutdown
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
defer shutdownCancel()
if err := srv.Shutdown(shutdownCtx); err != nil {
return fmt.Errorf("shutdown error: %v", err)
}
slog.Debug("Shutdown called, waiting for connections to close...")
// Wait for server to finish
select {
case err := <-serverErr:
// http.ErrServerClosed is expected when shutting down
if err != nil && !errors.Is(err, http.ErrServerClosed) {
return fmt.Errorf("server error during shutdown: %v", err)
}
case <-shutdownCtx.Done():
slog.Warn(fmt.Sprintf("%s server did not stop gracefully within timeout", transportName))
}
}
return nil
}
func handleHealthz(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte("ok"))
}
func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt disabledTools, gc mcpgrafana.GrafanaConfig, tls tlsConfig) error {
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: logLevel})))
s, tm := newServer(transport, dt)
// Create a context that will be cancelled on shutdown
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Set up signal handling for graceful shutdown
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
defer signal.Stop(sigChan)
// Handle shutdown signals
go func() {
<-sigChan
slog.Info("Received shutdown signal")
cancel()
// For stdio, close stdin to unblock the Listen call
if transport == "stdio" {
_ = os.Stdin.Close()
}
}()
// Start the appropriate server based on transport
switch transport {
case "stdio":
srv := server.NewStdioServer(s)
cf := mcpgrafana.ComposedStdioContextFunc(gc)
srv.SetContextFunc(cf)
// For stdio (single-tenant), initialize proxied tools on the server directly
if !dt.proxied {
stdioCtx := cf(ctx)
if err := tm.InitializeAndRegisterServerTools(stdioCtx); err != nil {
slog.Error("failed to initialize proxied tools for stdio", "error", err)
}
}
slog.Info("Starting Grafana MCP server using stdio transport", "version", mcpgrafana.Version())
err := srv.Listen(ctx, os.Stdin, os.Stdout)
if err != nil && err != context.Canceled {
return fmt.Errorf("server error: %v", err)
}
return nil
case "sse":
httpSrv := &http.Server{Addr: addr}
srv := server.NewSSEServer(s,
server.WithSSEContextFunc(mcpgrafana.ComposedSSEContextFunc(gc)),
server.WithStaticBasePath(basePath),
server.WithHTTPServer(httpSrv),
)
mux := http.NewServeMux()
if basePath == "" {
basePath = "/"
}
mux.Handle(basePath, srv)
mux.HandleFunc("/healthz", handleHealthz)
httpSrv.Handler = mux
slog.Info("Starting Grafana MCP server using SSE transport",
"version", mcpgrafana.Version(), "address", addr, "basePath", basePath)
return runHTTPServer(ctx, srv, addr, "SSE")
case "streamable-http":
httpSrv := &http.Server{Addr: addr}
opts := []server.StreamableHTTPOption{
server.WithHTTPContextFunc(mcpgrafana.ComposedHTTPContextFunc(gc)),
server.WithStateLess(dt.proxied), // Stateful when proxied tools enabled (requires sessions)
server.WithEndpointPath(endpointPath),
server.WithStreamableHTTPServer(httpSrv),
}
if tls.certFile != "" || tls.keyFile != "" {
opts = append(opts, server.WithTLSCert(tls.certFile, tls.keyFile))
}
srv := server.NewStreamableHTTPServer(s, opts...)
mux := http.NewServeMux()
mux.Handle(endpointPath, srv)
mux.HandleFunc("/healthz", handleHealthz)
httpSrv.Handler = mux
slog.Info("Starting Grafana MCP server using StreamableHTTP transport",
"version", mcpgrafana.Version(), "address", addr, "endpointPath", endpointPath)
return runHTTPServer(ctx, srv, addr, "StreamableHTTP")
default:
return fmt.Errorf("invalid transport type: %s. Must be 'stdio', 'sse' or 'streamable-http'", transport)
}
}
func main() {
var transport string
flag.StringVar(&transport, "t", "stdio", "Transport type (stdio, sse or streamable-http)")
flag.StringVar(
&transport,
"transport",
"stdio",
"Transport type (stdio, sse or streamable-http)",
)
addr := flag.String("address", "localhost:8000", "The host and port to start the sse server on")
basePath := flag.String("base-path", "", "Base path for the sse server")
endpointPath := flag.String("endpoint-path", "/mcp", "Endpoint path for the streamable-http server")
logLevel := flag.String("log-level", "info", "Log level (debug, info, warn, error)")
showVersion := flag.Bool("version", false, "Print the version and exit")
var dt disabledTools
dt.addFlags()
var gc grafanaConfig
gc.addFlags()
var tls tlsConfig
tls.addFlags()
flag.Parse()
if *showVersion {
fmt.Println(mcpgrafana.Version())
os.Exit(0)
}
// Convert local grafanaConfig to mcpgrafana.GrafanaConfig
grafanaConfig := mcpgrafana.GrafanaConfig{Debug: gc.debug}
if gc.tlsCertFile != "" || gc.tlsKeyFile != "" || gc.tlsCAFile != "" || gc.tlsSkipVerify {
grafanaConfig.TLSConfig = &mcpgrafana.TLSConfig{
CertFile: gc.tlsCertFile,
KeyFile: gc.tlsKeyFile,
CAFile: gc.tlsCAFile,
SkipVerify: gc.tlsSkipVerify,
}
}
if err := run(transport, *addr, *basePath, *endpointPath, parseLevel(*logLevel), dt, grafanaConfig, tls); err != nil {
panic(err)
}
}
func parseLevel(level string) slog.Level {
var l slog.Level
if err := l.UnmarshalText([]byte(level)); err != nil {
return slog.LevelInfo
}
return l
}
```
--------------------------------------------------------------------------------
/session_test.go:
--------------------------------------------------------------------------------
```go
//go:build integration
// Integration tests for proxied MCP tools functionality.
// Requires docker-compose to be running with Grafana and Tempo instances.
// Run with: go test -tags integration -v ./...
package mcpgrafana
import (
"context"
"fmt"
"net/url"
"os"
"strings"
"sync"
"testing"
"github.com/go-openapi/strfmt"
grafana_client "github.com/grafana/grafana-openapi-client-go/client"
"github.com/mark3labs/mcp-go/mcp"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// newProxiedToolsTestContext creates a test context with Grafana client and config
func newProxiedToolsTestContext(t *testing.T) context.Context {
cfg := grafana_client.DefaultTransportConfig()
cfg.Host = "localhost:3000"
cfg.Schemes = []string{"http"}
// Extract transport config from env vars, and set it on the context.
if u, ok := os.LookupEnv("GRAFANA_URL"); ok {
parsedURL, err := url.Parse(u)
require.NoError(t, err, "invalid GRAFANA_URL")
cfg.Host = parsedURL.Host
// The Grafana client will always prefer HTTPS even if the URL is HTTP,
// so we need to limit the schemes to HTTP if the URL is HTTP.
if parsedURL.Scheme == "http" {
cfg.Schemes = []string{"http"}
}
}
// Check for the new service account token environment variable first
if apiKey := os.Getenv("GRAFANA_SERVICE_ACCOUNT_TOKEN"); apiKey != "" {
cfg.APIKey = apiKey
} else if apiKey := os.Getenv("GRAFANA_API_KEY"); apiKey != "" {
// Fall back to the deprecated API key environment variable
cfg.APIKey = apiKey
} else {
cfg.BasicAuth = url.UserPassword("admin", "admin")
}
grafanaClient := grafana_client.NewHTTPClientWithConfig(strfmt.Default, cfg)
grafanaCfg := GrafanaConfig{
Debug: true,
URL: "http://localhost:3000",
APIKey: cfg.APIKey,
BasicAuth: cfg.BasicAuth,
}
ctx := WithGrafanaConfig(context.Background(), grafanaCfg)
return WithGrafanaClient(ctx, grafanaClient)
}
func TestDiscoverMCPDatasources(t *testing.T) {
ctx := newProxiedToolsTestContext(t)
t.Run("discovers tempo datasources", func(t *testing.T) {
discovered, err := discoverMCPDatasources(ctx)
require.NoError(t, err)
// Should find two Tempo datasources from docker-compose
assert.GreaterOrEqual(t, len(discovered), 2, "Should discover at least 2 Tempo datasources")
// Check that we found the expected datasources
uids := make([]string, len(discovered))
for i, ds := range discovered {
uids[i] = ds.UID
assert.Equal(t, "tempo", ds.Type, "All discovered datasources should be tempo type")
assert.NotEmpty(t, ds.Name, "Datasource should have a name")
assert.NotEmpty(t, ds.MCPURL, "Datasource should have MCP URL")
// Verify URL format
expectedURLPattern := fmt.Sprintf("http://localhost:3000/api/datasources/proxy/uid/%s/api/mcp", ds.UID)
assert.Equal(t, expectedURLPattern, ds.MCPURL, "MCP URL should follow proxy pattern")
}
// Should contain our expected UIDs
assert.Contains(t, uids, "tempo", "Should discover 'tempo' datasource")
assert.Contains(t, uids, "tempo-secondary", "Should discover 'tempo-secondary' datasource")
})
t.Run("returns error when grafana client not in context", func(t *testing.T) {
emptyCtx := context.Background()
discovered, err := discoverMCPDatasources(emptyCtx)
assert.Error(t, err)
assert.Nil(t, discovered)
assert.Contains(t, err.Error(), "grafana client not found in context")
})
t.Run("returns error when auth is missing", func(t *testing.T) {
// Context with client but no auth credentials
cfg := grafana_client.DefaultTransportConfig()
cfg.Host = "localhost:3000"
cfg.Schemes = []string{"http"}
grafanaClient := grafana_client.NewHTTPClientWithConfig(strfmt.Default, cfg)
grafanaCfg := GrafanaConfig{
URL: "http://localhost:3000",
// No APIKey or BasicAuth set
}
ctx := WithGrafanaConfig(context.Background(), grafanaCfg)
ctx = WithGrafanaClient(ctx, grafanaClient)
discovered, err := discoverMCPDatasources(ctx)
assert.Error(t, err)
assert.Nil(t, discovered)
assert.Contains(t, err.Error(), "Unauthorized")
})
}
func TestToolNamespacing(t *testing.T) {
t.Run("parse proxied tool name", func(t *testing.T) {
datasourceType, toolName, err := parseProxiedToolName("tempo_traceql-search")
require.NoError(t, err)
assert.Equal(t, "tempo", datasourceType)
assert.Equal(t, "traceql-search", toolName)
})
t.Run("parse proxied tool name with multiple underscores", func(t *testing.T) {
datasourceType, toolName, err := parseProxiedToolName("tempo_get-attribute-values")
require.NoError(t, err)
assert.Equal(t, "tempo", datasourceType)
assert.Equal(t, "get-attribute-values", toolName)
})
t.Run("parse proxied tool name with invalid format", func(t *testing.T) {
_, _, err := parseProxiedToolName("invalid")
assert.Error(t, err)
assert.Contains(t, err.Error(), "invalid proxied tool name format")
})
t.Run("add datasourceUid parameter to tool", func(t *testing.T) {
originalTool := mcp.Tool{
Name: "query_traces",
Description: "Query traces from Tempo",
InputSchema: mcp.ToolInputSchema{
Properties: map[string]any{
"query": map[string]any{
"type": "string",
},
},
Required: []string{"query"},
},
}
modifiedTool := addDatasourceUidParameter(originalTool, "tempo")
assert.Equal(t, "tempo_query_traces", modifiedTool.Name)
assert.Equal(t, "Query traces from Tempo", modifiedTool.Description)
assert.NotNil(t, modifiedTool.InputSchema.Properties["datasourceUid"])
assert.Contains(t, modifiedTool.InputSchema.Required, "datasourceUid")
assert.Contains(t, modifiedTool.InputSchema.Required, "query")
})
t.Run("add datasourceUid parameter with empty description", func(t *testing.T) {
originalTool := mcp.Tool{
Name: "test_tool",
Description: "",
InputSchema: mcp.ToolInputSchema{
Properties: make(map[string]any),
},
}
modifiedTool := addDatasourceUidParameter(originalTool, "tempo")
assert.Equal(t, "tempo_test_tool", modifiedTool.Name)
assert.Equal(t, "", modifiedTool.Description, "Should not modify empty description")
assert.NotNil(t, modifiedTool.InputSchema.Properties["datasourceUid"])
})
}
func TestSessionStateLifecycle(t *testing.T) {
t.Run("create and get session", func(t *testing.T) {
sm := NewSessionManager()
// Create mock session
mockSession := &mockClientSession{id: "test-session-123"}
sm.CreateSession(context.Background(), mockSession)
state, exists := sm.GetSession("test-session-123")
assert.True(t, exists)
assert.NotNil(t, state)
assert.NotNil(t, state.proxiedClients)
assert.False(t, state.proxiedToolsInitialized)
})
t.Run("remove session cleans up clients", func(t *testing.T) {
sm := NewSessionManager()
mockSession := &mockClientSession{id: "test-session-456"}
sm.CreateSession(context.Background(), mockSession)
state, _ := sm.GetSession("test-session-456")
// Add a mock proxied client
mockClient := &ProxiedClient{
DatasourceUID: "test-uid",
DatasourceName: "Test Datasource",
DatasourceType: "tempo",
}
state.proxiedClients["tempo_test-uid"] = mockClient
// Remove session
sm.RemoveSession(context.Background(), mockSession)
// Session should be gone
_, exists := sm.GetSession("test-session-456")
assert.False(t, exists)
})
t.Run("get non-existent session", func(t *testing.T) {
sm := NewSessionManager()
state, exists := sm.GetSession("non-existent")
assert.False(t, exists)
assert.Nil(t, state)
})
}
func TestConcurrentInitializationRaceCondition(t *testing.T) {
t.Run("concurrent initialization calls should be safe", func(t *testing.T) {
sm := NewSessionManager()
mockSession := &mockClientSession{id: "race-test-session"}
sm.CreateSession(context.Background(), mockSession)
state, exists := sm.GetSession("race-test-session")
require.True(t, exists)
// Track how many times the initialization logic runs
var initCount int
var initCountMutex sync.Mutex
// Create a custom initOnce to track calls
state.initOnce = sync.Once{}
// Simulate the initialization work that should run exactly once
initWork := func() {
initCountMutex.Lock()
initCount++
initCountMutex.Unlock()
// Simulate some work
state.mutex.Lock()
state.proxiedToolsInitialized = true
state.proxiedClients["tempo_test"] = &ProxiedClient{
DatasourceUID: "test",
DatasourceName: "Test",
DatasourceType: "tempo",
}
state.mutex.Unlock()
}
// Launch multiple goroutines that all try to initialize concurrently
const numGoroutines = 10
var wg sync.WaitGroup
wg.Add(numGoroutines)
for i := 0; i < numGoroutines; i++ {
go func() {
defer wg.Done()
// This should be the pattern used in InitializeAndRegisterProxiedTools
state.initOnce.Do(initWork)
}()
}
wg.Wait()
// Verify initialization ran exactly once
assert.Equal(t, 1, initCount, "Initialization should run exactly once despite concurrent calls")
assert.True(t, state.proxiedToolsInitialized, "State should be initialized")
assert.Len(t, state.proxiedClients, 1, "Should have exactly one client")
})
t.Run("sync.Once prevents double initialization", func(t *testing.T) {
sm := NewSessionManager()
mockSession := &mockClientSession{id: "double-init-test"}
sm.CreateSession(context.Background(), mockSession)
state, _ := sm.GetSession("double-init-test")
callCount := 0
// First call
state.initOnce.Do(func() {
callCount++
})
// Second call should not execute
state.initOnce.Do(func() {
callCount++
})
// Third call should also not execute
state.initOnce.Do(func() {
callCount++
})
assert.Equal(t, 1, callCount, "sync.Once should ensure function runs exactly once")
})
}
func TestProxiedClientLifecycle(t *testing.T) {
ctx := newProxiedToolsTestContext(t)
t.Run("list tools returns copy", func(t *testing.T) {
pc := &ProxiedClient{
DatasourceUID: "test-uid",
DatasourceName: "Test",
DatasourceType: "tempo",
Tools: []mcp.Tool{
{Name: "tool1", Description: "First tool"},
{Name: "tool2", Description: "Second tool"},
},
}
tools1 := pc.ListTools()
tools2 := pc.ListTools()
// Should return same content
assert.Equal(t, tools1, tools2)
// But different slice instances (copy)
assert.NotSame(t, &tools1[0], &tools2[0])
})
t.Run("call tool validates tool exists", func(t *testing.T) {
pc := &ProxiedClient{
DatasourceUID: "test-uid",
DatasourceName: "Test",
DatasourceType: "tempo",
Tools: []mcp.Tool{
{Name: "valid_tool", Description: "Valid tool"},
},
}
// Call non-existent tool
result, err := pc.CallTool(ctx, "non_existent_tool", map[string]any{})
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "not found in remote MCP server")
})
}
func TestEndToEndProxiedToolsFlow(t *testing.T) {
ctx := newProxiedToolsTestContext(t)
t.Run("full flow from discovery to tool call", func(t *testing.T) {
// Step 1: Discover MCP datasources
discovered, err := discoverMCPDatasources(ctx)
require.NoError(t, err)
require.GreaterOrEqual(t, len(discovered), 1, "Should discover at least one Tempo datasource")
// Use the first discovered datasource
ds := discovered[0]
t.Logf("Testing with datasource: %s (UID: %s, URL: %s)", ds.Name, ds.UID, ds.MCPURL)
// Step 2: Create a proxied client connection
client, err := NewProxiedClient(ctx, ds.UID, ds.Name, ds.Type, ds.MCPURL)
if err != nil {
t.Skipf("Skipping end-to-end test: Tempo MCP endpoint not available: %v", err)
return
}
defer func() {
_ = client.Close()
}()
// Step 3: Verify we got tools from the remote server
tools := client.ListTools()
require.Greater(t, len(tools), 0, "Should have at least one tool from Tempo MCP server")
t.Logf("Discovered %d tools from Tempo MCP server", len(tools))
// Log the available tools
for _, tool := range tools {
t.Logf(" - Tool: %s - %s", tool.Name, tool.Description)
}
// Step 4: Test tool modification with datasourceUid parameter
firstTool := tools[0]
modifiedTool := addDatasourceUidParameter(firstTool, ds.Type)
expectedName := ds.Type + "_" + firstTool.Name
assert.Equal(t, expectedName, modifiedTool.Name, "Modified tool should have prefixed name")
assert.Contains(t, modifiedTool.InputSchema.Required, "datasourceUid", "Modified tool should require datasourceUid")
// Step 5: Test session integration
sm := NewSessionManager()
mockSession := &mockClientSession{id: "e2e-test-session"}
sm.CreateSession(ctx, mockSession)
state, exists := sm.GetSession("e2e-test-session")
require.True(t, exists)
// Store the proxied client in session state
key := ds.Type + "_" + ds.UID
state.proxiedClients[key] = client
// Step 6: Verify client is stored correctly in session
retrievedClient, exists := state.proxiedClients[key]
require.True(t, exists, "Client should be stored in session state")
assert.Equal(t, client, retrievedClient, "Should retrieve the same client from session")
// Step 7: Test ProxiedToolHandler flow
handler := NewProxiedToolHandler(sm, nil, modifiedTool.Name)
assert.NotNil(t, handler)
// Note: We can't actually call the tool without knowing what arguments it expects
// and without the context having the proper session, but we've validated the setup
t.Logf("Successfully validated end-to-end proxied tools flow")
})
t.Run("multiple datasources in single session", func(t *testing.T) {
discovered, err := discoverMCPDatasources(ctx)
require.NoError(t, err)
if len(discovered) < 2 {
t.Skip("Need at least 2 Tempo datasources for this test")
}
sm := NewSessionManager()
mockSession := &mockClientSession{id: "multi-ds-test-session"}
sm.CreateSession(ctx, mockSession)
state, _ := sm.GetSession("multi-ds-test-session")
// Try to connect to multiple datasources
connectedCount := 0
for i, ds := range discovered {
if i >= 2 {
break // Test with first 2 datasources
}
client, err := NewProxiedClient(ctx, ds.UID, ds.Name, ds.Type, ds.MCPURL)
if err != nil {
t.Logf("Could not connect to datasource %s: %v", ds.UID, err)
continue
}
defer func() {
_ = client.Close()
}()
key := ds.Type + "_" + ds.UID
state.proxiedClients[key] = client
connectedCount++
t.Logf("Connected to datasource %s with %d tools", ds.UID, len(client.Tools))
}
if connectedCount == 0 {
t.Skip("Could not connect to any Tempo datasources")
}
// Verify each client is stored correctly
for key, client := range state.proxiedClients {
parts := strings.Split(key, "_")
require.Len(t, parts, 2, "Key should have format type_uid")
assert.NotNil(t, client, "Client should not be nil")
assert.Equal(t, parts[0], client.DatasourceType, "Client type should match key")
assert.Equal(t, parts[1], client.DatasourceUID, "Client UID should match key")
}
t.Logf("Successfully managed %d datasources in single session", connectedCount)
})
}
```