feat: add support for proxying to Tempo datasources' MCP servers

Makefile

@@ -53,15 +53,15 @@ test-python-e2e: ## Run Python E2E tests (requires docker-compose services and S
 
 .PHONY: run
 run: ## Run the MCP server in stdio mode.
-	go run ./cmd/mcp-grafana
+	GRAFANA_USERNAME=admin GRAFANA_PASSWORD=admin go run ./cmd/mcp-grafana
 
 .PHONY: run-sse
 run-sse: ## Run the MCP server in SSE mode.
-	go run ./cmd/mcp-grafana --transport sse --log-level debug --debug
+	GRAFANA_USERNAME=admin GRAFANA_PASSWORD=admin go run ./cmd/mcp-grafana --transport sse --log-level debug --debug
 
 PHONY: run-streamable-http
 run-streamable-http: ## Run the MCP server in StreamableHTTP mode.
-	go run ./cmd/mcp-grafana --transport streamable-http --log-level debug --debug
+	GRAFANA_USERNAME=admin GRAFANA_PASSWORD=admin go run ./cmd/mcp-grafana --transport streamable-http --log-level debug --debug
 
 .PHONY: run-test-services
 run-test-services: ## Run the docker-compose services required for the unit and integration tests.

cmd/mcp-grafana/main.go

@@ -14,6 +14,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
 
 	mcpgrafana "github.com/grafana/mcp-grafana"
@@ -40,7 +41,7 @@ type disabledTools struct {
 	search, datasource, incident,
 	prometheus, loki, alerting,
 	dashboard, folder, oncall, asserts, sift, admin,
-	pyroscope, navigation bool
+	pyroscope, navigation, proxied bool
 }
 
 // Configuration for the Grafana client.
@@ -56,8 +57,7 @@ type grafanaConfig struct {
 }
 
 func (dt *disabledTools) addFlags() {
-	flag.StringVar(&dt.enabledTools, "enabled-tools", "search,datasource,incident,prometheus,loki,alerting,dashboard,folder,oncall,asserts,sift,admin,pyroscope,navigation", "A comma separated list of tools enabled for this server. Can be overwritten entirely or by disabling specific components, e.g. --disable-search.")
-
+	flag.StringVar(&dt.enabledTools, "enabled-tools", "search,datasource,incident,prometheus,loki,alerting,dashboard,folder,oncall,asserts,sift,admin,pyroscope,navigation,proxied", "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")
@@ -72,6 +72,7 @@ func (dt *disabledTools) addFlags() {
 	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)")
 }
 
 func (gc *grafanaConfig) addFlags() {
@@ -102,24 +103,71 @@ func (dt *disabledTools) addTools(s *server.MCPServer) {
 	maybeAddTools(s, tools.AddNavigationTools, enabledTools, dt.navigation, "navigation")
 }
 
-func newServer(dt disabledTools) *server.MCPServer {
-	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.
-	`))
+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
+	return s, stm
 }
 
 type tlsConfig struct {
@@ -162,6 +210,7 @@ func runHTTPServer(ctx context.Context, srv httpServer, addr, transportName stri
 		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 {
@@ -180,7 +229,7 @@ func runHTTPServer(ctx context.Context, srv httpServer, addr, transportName stri
 
 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 := newServer(dt)
+	s, tm := newServer(transport, dt)
 
 	// Create a context that will be cancelled on shutdown
 	ctx, cancel := context.WithCancel(context.Background())
@@ -207,7 +256,17 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 	switch transport {
 	case "stdio":
 		srv := server.NewStdioServer(s)
-		srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(gc))
+		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)
@@ -227,7 +286,7 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 	case "streamable-http":
 		opts := []server.StreamableHTTPOption{
 			server.WithHTTPContextFunc(mcpgrafana.ComposedHTTPContextFunc(gc)),
-			server.WithStateLess(true),
+			server.WithStateLess(dt.proxied), // Stateful when proxied tools enabled (requires sessions)
 			server.WithEndpointPath(endpointPath),
 		}
 		if tls.certFile != "" || tls.keyFile != "" {
@@ -238,10 +297,7 @@ func run(transport, addr, basePath, endpointPath string, logLevel slog.Level, dt
 			"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,
-		)
+		return fmt.Errorf("invalid transport type: %s. Must be 'stdio', 'sse' or 'streamable-http'", transport)
 	}
 }
 

docker-compose.yaml

@@ -43,3 +43,19 @@ services:
     image: grafana/pyroscope:1.13.4
     ports:
       - 4040:4040
+
+  tempo:
+    image: grafana/tempo:2.9.0-rc.0
+    command: ["-config.file=/etc/tempo/tempo-config.yaml"]
+    volumes:
+      - ./testdata/tempo-config.yaml:/etc/tempo/tempo-config.yaml
+    ports:
+      - "3200:3200" # tempo
+
+  tempo2:
+    image: grafana/tempo:2.9.0-rc.0
+    command: ["-config.file=/etc/tempo/tempo-config.yaml"]
+    volumes:
+      - ./testdata/tempo-config-2.yaml:/etc/tempo/tempo-config.yaml
+    ports:
+      - "3201:3201" # tempo instance 2

go.mod

@@ -6,6 +6,7 @@ require (
 	connectrpc.com/connect v1.19.0
 	github.com/PaesslerAG/gval v1.2.4
 	github.com/PaesslerAG/jsonpath v0.1.1
+	github.com/cenkalti/backoff/v5 v5.0.3
 	github.com/go-openapi/runtime v0.29.0
 	github.com/go-openapi/strfmt v0.24.0
 	github.com/google/uuid v1.6.0
@@ -31,7 +32,6 @@ require (
 	github.com/bahlo/generic-list-go v0.2.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
-	github.com/cenkalti/backoff/v5 v5.0.3 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/cheekybits/genny v1.0.0 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect

mcpgrafana.go

@@ -435,14 +435,15 @@ var ExtractGrafanaClientFromEnv server.StdioContextFunc = func(ctx context.Conte
 	}
 	auth := userAndPassFromEnv()
 	grafanaClient := NewGrafanaClient(ctx, grafanaURL, apiKey, auth)
-	return context.WithValue(ctx, grafanaClientKey{}, grafanaClient)
+	return WithGrafanaClient(ctx, grafanaClient)
 }
 
 // ExtractGrafanaClientFromHeaders is a HTTPContextFunc that creates and injects a Grafana client into the context.
 // It prioritizes configuration from HTTP headers (X-Grafana-URL, X-Grafana-API-Key) over environment variables for multi-tenant scenarios.
 var ExtractGrafanaClientFromHeaders httpContextFunc = func(ctx context.Context, req *http.Request) context.Context {
 	// Extract transport config from request headers, and set it on the context.
 	u, apiKey, basicAuth := extractKeyGrafanaInfoFromReq(req)
+	slog.Debug("Creating Grafana client", "url", u, "api_key_set", apiKey != "", "basic_auth_set", basicAuth != nil)
 
 	grafanaClient := NewGrafanaClient(ctx, u, apiKey, basicAuth)
 	return WithGrafanaClient(ctx, grafanaClient)

proxied_client.go

@@ -0,0 +1,143 @@
+package mcpgrafana
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"log/slog"
+	"sync"
+
+	mcp_client "github.com/mark3labs/mcp-go/client"
+	"github.com/mark3labs/mcp-go/client/transport"
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// ProxiedClient represents a connection to a remote MCP server (e.g., Tempo datasource)
+type ProxiedClient struct {
+	DatasourceUID  string
+	DatasourceName string
+	DatasourceType string
+	Client         *mcp_client.Client
+	Tools          []mcp.Tool
+	mutex          sync.RWMutex
+}
+
+// NewProxiedClient creates a new connection to a remote MCP server
+func NewProxiedClient(ctx context.Context, datasourceUID, datasourceName, datasourceType, mcpEndpoint string) (*ProxiedClient, error) {
+	// Get Grafana config for authentication
+	config := GrafanaConfigFromContext(ctx)
+
+	// Build headers for authentication
+	headers := make(map[string]string)
+	if config.APIKey != "" {
+		headers["Authorization"] = "Bearer " + config.APIKey
+	} else if config.BasicAuth != nil {
+		auth := config.BasicAuth.String()
+		headers["Authorization"] = "Basic " + base64.StdEncoding.EncodeToString([]byte(auth))
+	}
+
+	// Create HTTP transport with authentication headers
+	slog.DebugContext(ctx, "connecting to MCP server", "datasource", datasourceUID, "url", mcpEndpoint)
+	httpTransport, err := transport.NewStreamableHTTP(
+		mcpEndpoint,
+		transport.WithHTTPHeaders(headers),
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create HTTP transport: %w", err)
+	}
+
+	// Create MCP client
+	mcpClient := mcp_client.NewClient(httpTransport)
+
+	// Initialize the connection
+	initReq := mcp.InitializeRequest{}
+	initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
+	initReq.Params.ClientInfo = mcp.Implementation{
+		Name:    "mcp-grafana-proxy",
+		Version: Version(),
+	}
+
+	_, err = mcpClient.Initialize(ctx, initReq)
+	if err != nil {
+		_ = mcpClient.Close()
+		return nil, fmt.Errorf("failed to initialize MCP client: %w", err)
+	}
+
+	// List available tools from the remote server
+	listReq := mcp.ListToolsRequest{}
+	toolsResult, err := mcpClient.ListTools(ctx, listReq)
+	if err != nil {
+		_ = mcpClient.Close()
+		return nil, fmt.Errorf("failed to list tools from remote MCP server: %w", err)
+	}
+
+	slog.DebugContext(ctx, "connected to proxied MCP server",
+		"datasource", datasourceUID,
+		"type", datasourceType,
+		"tools", len(toolsResult.Tools))
+
+	return &ProxiedClient{
+		DatasourceUID:  datasourceUID,
+		DatasourceName: datasourceName,
+		DatasourceType: datasourceType,
+		Client:         mcpClient,
+		Tools:          toolsResult.Tools,
+	}, nil
+}
+
+// CallTool forwards a tool call to the remote MCP server
+func (pc *ProxiedClient) CallTool(ctx context.Context, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) {
+	pc.mutex.RLock()
+	defer pc.mutex.RUnlock()
+
+	// Validate the tool exists
+	var toolExists bool
+	for _, tool := range pc.Tools {
+		if tool.Name == toolName {
+			toolExists = true
+			break
+		}
+	}
+	if !toolExists {
+		return nil, fmt.Errorf("tool %s not found in remote MCP server", toolName)
+	}
+
+	// Create the call tool request
+	req := mcp.CallToolRequest{}
+	req.Params.Name = toolName
+	req.Params.Arguments = arguments
+
+	// Forward the call to the remote server
+	result, err := pc.Client.CallTool(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to call tool on remote MCP server: %w", err)
+	}
+
+	return result, nil
+}
+
+// ListTools returns the tools available from this remote server
+// Note: This method doesn't take a context parameter as the tools are cached locally
+func (pc *ProxiedClient) ListTools() []mcp.Tool {
+	pc.mutex.RLock()
+	defer pc.mutex.RUnlock()
+
+	// Return a copy to prevent external modification
+	result := make([]mcp.Tool, len(pc.Tools))
+	copy(result, pc.Tools)
+	return result
+}
+
+// Close closes the connection to the remote MCP server
+func (pc *ProxiedClient) Close() error {
+	pc.mutex.Lock()
+	defer pc.mutex.Unlock()
+
+	if pc.Client != nil {
+		if err := pc.Client.Close(); err != nil {
+			return fmt.Errorf("failed to close MCP client: %w", err)
+		}
+	}
+
+	return nil
+}