Add user agent tracking for minimal usage tracking (#229)

Author: sd2k

cmd/mcp-grafana/main.go

@@ -6,30 +6,15 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
-	"runtime/debug"
 	"slices"
 	"strings"
-	"sync"
 
 	"github.com/mark3labs/mcp-go/server"
 
 	mcpgrafana "github.com/grafana/mcp-grafana"
 	"github.com/grafana/mcp-grafana/tools"
 )
 
-// version returns the version of the mcp-grafana binary.
-// It is populated by the `runtime/debug` package which
-// fetches git information from the build directory.
-var version = sync.OnceValue(func() string {
-	// Default version string returned by `runtime/debug` if built
-	// from the source repository rather than with `go install`.
-	v := "(devel)"
-	if bi, ok := debug.ReadBuildInfo(); ok {
-		v = bi.Main.Version
-	}
-	return v
-})
-
 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)
@@ -111,7 +96,7 @@ func (dt *disabledTools) addTools(s *server.MCPServer) {
 }
 
 func newServer(dt disabledTools) *server.MCPServer {
-	s := server.NewMCPServer("mcp-grafana", version(), server.WithInstructions(`
+	s := server.NewMCPServer("mcp-grafana", mcpgrafana.Version(), server.WithInstructions(`
 	This server provides access to your Grafana instance and the surrounding ecosystem.
 
 	Available Capabilities:
@@ -138,14 +123,14 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 	case "stdio":
 		srv := server.NewStdioServer(s)
 		srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(gc))
-		slog.Info("Starting Grafana MCP server using stdio transport", "version", version())
+		slog.Info("Starting Grafana MCP server using stdio transport", "version", mcpgrafana.Version())
 		return srv.Listen(context.Background(), os.Stdin, os.Stdout)
 	case "sse":
 		srv := server.NewSSEServer(s,
 			server.WithSSEContextFunc(mcpgrafana.ComposedSSEContextFunc(gc)),
 			server.WithStaticBasePath(basePath),
 		)
-		slog.Info("Starting Grafana MCP server using SSE transport", "version", version(), "address", addr, "basePath", basePath)
+		slog.Info("Starting Grafana MCP server using SSE transport", "version", mcpgrafana.Version(), "address", addr, "basePath", basePath)
 		if err := srv.Start(addr); err != nil {
 			return fmt.Errorf("server error: %v", err)
 		}
@@ -154,7 +139,7 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 			server.WithStateLess(true),
 			server.WithEndpointPath(endpointPath),
 		)
-		slog.Info("Starting Grafana MCP server using StreamableHTTP transport", "version", version(), "address", addr, "endpointPath", endpointPath)
+		slog.Info("Starting Grafana MCP server using StreamableHTTP transport", "version", mcpgrafana.Version(), "address", addr, "endpointPath", endpointPath)
 		if err := srv.Start(addr); err != nil {
 			return fmt.Errorf("server error: %v", err)
 		}
@@ -188,7 +173,7 @@ func main() {
 	flag.Parse()
 
 	if *showVersion {
-		fmt.Println(version())
+		fmt.Println(mcpgrafana.Version())
 		os.Exit(0)
 	}
 

mcpgrafana.go

@@ -10,7 +10,9 @@ import (
 	"net/url"
 	"os"
 	"reflect"
+	"runtime/debug"
 	"strings"
+	"sync"
 
 	"github.com/go-openapi/strfmt"
 	"github.com/grafana/grafana-openapi-client-go/client"
@@ -59,7 +61,7 @@ type GrafanaConfig struct {
 	Debug bool
 
 	// IncludeArgumentsInSpans enables logging of tool arguments in OpenTelemetry spans.
-	// This should only be enabled in non-production environments or when you're certain 
+	// This should only be enabled in non-production environments or when you're certain
 	// the arguments don't contain PII. Defaults to false for safety.
 	// Note: OpenTelemetry spans are always created for context propagation, but arguments
 	// are only included when this flag is enabled.
@@ -147,6 +149,65 @@ func (tc *TLSConfig) HTTPTransport(defaultTransport *http.Transport) (http.Round
 	return transport, nil
 }
 
+// UserAgentTransport wraps an http.RoundTripper to add a custom User-Agent header
+type UserAgentTransport struct {
+	rt        http.RoundTripper
+	UserAgent string
+}
+
+func (t *UserAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	// Clone the request to avoid modifying the original
+	clonedReq := req.Clone(req.Context())
+
+	// Add or update the User-Agent header
+	if clonedReq.Header.Get("User-Agent") == "" {
+		clonedReq.Header.Set("User-Agent", t.UserAgent)
+	}
+
+	return t.rt.RoundTrip(clonedReq)
+}
+
+// Version returns the version of the mcp-grafana binary.
+// It is populated by the `runtime/debug` package which
+// fetches git information from the build directory.
+var Version = sync.OnceValue(func() string {
+	// Default version string returned by `runtime/debug` if built
+	// from the source repository rather than with `go install`.
+	v := "(devel)"
+	if bi, ok := debug.ReadBuildInfo(); ok && bi.Main.Version != "" {
+		v = bi.Main.Version
+	}
+	return v
+})
+
+// UserAgent returns the user agent string for HTTP requests
+func UserAgent() string {
+	return fmt.Sprintf("mcp-grafana/%s", Version())
+}
+
+// NewUserAgentTransport creates a new UserAgentTransport with the specified user agent.
+// If no user agent is provided, uses the default UserAgent().
+func NewUserAgentTransport(rt http.RoundTripper, userAgent ...string) *UserAgentTransport {
+	if rt == nil {
+		rt = http.DefaultTransport
+	}
+
+	ua := UserAgent() // default
+	if len(userAgent) > 0 {
+		ua = userAgent[0]
+	}
+
+	return &UserAgentTransport{
+		rt:        rt,
+		UserAgent: ua,
+	}
+}
+
+// wrapWithUserAgent wraps an http.RoundTripper with user agent tracking
+func wrapWithUserAgent(rt http.RoundTripper) http.RoundTripper {
+	return NewUserAgentTransport(rt)
+}
+
 // ExtractGrafanaInfoFromEnv is a StdioContextFunc that extracts Grafana configuration
 // from environment variables and injects a configured client into the context.
 var ExtractGrafanaInfoFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
@@ -281,9 +342,11 @@ func NewGrafanaClient(ctx context.Context, grafanaURL, apiKey string) *client.Gr
 			transportField := v.FieldByName("Transport")
 			if transportField.IsValid() && transportField.CanSet() {
 				if rt, ok := transportField.Interface().(http.RoundTripper); ok {
-					wrapped := otelhttp.NewTransport(rt)
+					// Wrap with user agent first, then otel
+					userAgentWrapped := wrapWithUserAgent(rt)
+					wrapped := otelhttp.NewTransport(userAgentWrapped)
 					transportField.Set(reflect.ValueOf(wrapped))
-					slog.Debug("HTTP tracing enabled for Grafana client")
+					slog.Debug("HTTP tracing and user agent tracking enabled for Grafana client")
 				}
 			}
 		}
@@ -363,12 +426,15 @@ var ExtractIncidentClientFromEnv server.StdioContextFunc = func(ctx context.Cont
 		if err != nil {
 			slog.Error("Failed to create custom transport for incident client, using default", "error", err)
 		} else {
-			client.HTTPClient.Transport = transport
-			slog.Debug("Using custom TLS configuration for incident client",
+			client.HTTPClient.Transport = wrapWithUserAgent(transport)
+			slog.Debug("Using custom TLS configuration and user agent for incident client",
 				"cert_file", tlsConfig.CertFile,
 				"ca_file", tlsConfig.CAFile,
 				"skip_verify", tlsConfig.SkipVerify)
 		}
+	} else {
+		// No custom TLS, but still add user agent
+		client.HTTPClient.Transport = wrapWithUserAgent(http.DefaultTransport)
 	}
 
 	return context.WithValue(ctx, incidentClientKey{}, client)
@@ -395,12 +461,15 @@ var ExtractIncidentClientFromHeaders httpContextFunc = func(ctx context.Context,
 		if err != nil {
 			slog.Error("Failed to create custom transport for incident client, using default", "error", err)
 		} else {
-			client.HTTPClient.Transport = transport
-			slog.Debug("Using custom TLS configuration for incident client",
+			client.HTTPClient.Transport = wrapWithUserAgent(transport)
+			slog.Debug("Using custom TLS configuration and user agent for incident client",
 				"cert_file", tlsConfig.CertFile,
 				"ca_file", tlsConfig.CAFile,
 				"skip_verify", tlsConfig.SkipVerify)
 		}
+	} else {
+		// No custom TLS, but still add user agent
+		client.HTTPClient.Transport = wrapWithUserAgent(http.DefaultTransport)
 	}
 
 	return context.WithValue(ctx, incidentClientKey{}, client)

mcpgrafana_test.go

@@ -80,7 +80,7 @@ func TestExtractGrafanaInfoFromHeaders(t *testing.T) {
 		// Explicitly clear environment variables to ensure test isolation
 		t.Setenv("GRAFANA_URL", "")
 		t.Setenv("GRAFANA_API_KEY", "")
-		
+
 		req, err := http.NewRequest("GET", "http://example.com", nil)
 		require.NoError(t, err)
 		ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
@@ -239,7 +239,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		type TestParams struct {
 			Message string `json:"message" jsonschema:"description=Test message"`
 		}
-		
+
 		testHandler := func(ctx context.Context, args TestParams) (string, error) {
 			return "Hello " + args.Message, nil
 		}
@@ -295,7 +295,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		type TestParams struct {
 			ShouldFail bool `json:"shouldFail" jsonschema:"description=Whether to fail"`
 		}
-		
+
 		testHandler := func(ctx context.Context, args TestParams) (string, error) {
 			if args.ShouldFail {
 				return "", assert.AnError
@@ -358,7 +358,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		type TestParams struct {
 			Message string `json:"message" jsonschema:"description=Test message"`
 		}
-		
+
 		testHandler := func(ctx context.Context, args TestParams) (string, error) {
 			return "processed", nil
 		}
@@ -406,7 +406,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		type TestParams struct {
 			SensitiveData string `json:"sensitiveData" jsonschema:"description=Potentially sensitive data"`
 		}
-		
+
 		testHandler := func(ctx context.Context, args TestParams) (string, error) {
 			return "processed", nil
 		}
@@ -451,7 +451,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		attributes := span.Attributes()
 		assertHasAttribute(t, attributes, "mcp.tool.name", "sensitive_tool")
 		assertHasAttribute(t, attributes, "mcp.tool.description", "A tool with sensitive data")
-		
+
 		// Verify arguments are NOT present
 		for _, attr := range attributes {
 			assert.NotEqual(t, "mcp.tool.arguments", string(attr.Key), "Arguments should not be logged by default for PII safety")
@@ -466,7 +466,7 @@ func TestToolTracingInstrumentation(t *testing.T) {
 		type TestParams struct {
 			SafeData string `json:"safeData" jsonschema:"description=Non-sensitive data"`
 		}
-		
+
 		testHandler := func(ctx context.Context, args TestParams) (string, error) {
 			return "processed", nil
 		}
@@ -531,7 +531,7 @@ func TestHTTPTracingConfiguration(t *testing.T) {
 
 	t.Run("tracing works gracefully without OpenTelemetry configured", func(t *testing.T) {
 		// No OpenTelemetry tracer provider configured
-		
+
 		// Create context (tracing always enabled for context propagation)
 		config := GrafanaConfig{}
 		ctx := WithGrafanaConfig(context.Background(), config)