fix: fall back to Grafana URL/API key from env vars in SSE mode (#94)

If there are no relevant headers, fall back to the Grafana URL and API key
from the environment variables, then the defaults.

Fixes #91.

Author: sd2k

mcpgrafana.go

@@ -146,6 +146,13 @@ var ExtractGrafanaClientFromHeaders server.SSEContextFunc = func(ctx context.Con
 	cfg := client.DefaultTransportConfig()
 	// Extract transport config from request headers, and set it on the context.
 	u, apiKey := urlAndAPIKeyFromHeaders(req)
+	uEnv, apiKeyEnv := urlAndAPIKeyFromEnv()
+	if u == "" {
+		u = uEnv
+	}
+	if apiKey == "" {
+		apiKey = apiKeyEnv
+	}
 	if u != "" {
 		if url, err := url.Parse(u); err == nil {
 			cfg.Host = url.Host
@@ -197,6 +204,16 @@ var ExtractIncidentClientFromEnv server.StdioContextFunc = func(ctx context.Cont
 
 var ExtractIncidentClientFromHeaders server.SSEContextFunc = func(ctx context.Context, req *http.Request) context.Context {
 	grafanaURL, apiKey := urlAndAPIKeyFromHeaders(req)
+	grafanaURLEnv, apiKeyEnv := urlAndAPIKeyFromEnv()
+	if grafanaURL == "" {
+		grafanaURL = grafanaURLEnv
+	}
+	if grafanaURL == "" {
+		grafanaURL = defaultGrafanaURL
+	}
+	if apiKey == "" {
+		apiKey = apiKeyEnv
+	}
 	incidentURL := fmt.Sprintf("%s/api/plugins/grafana-irm-app/resources/api/v1/", grafanaURL)
 	client := incident.NewClient(incidentURL, apiKey)
 	return context.WithValue(ctx, incidentClientKey{}, client)

mcpgrafana_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/go-openapi/runtime/client"
+	grafana_client "github.com/grafana/grafana-openapi-client-go/client"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -22,6 +23,52 @@ func TestExtractIncidentClientFromEnv(t *testing.T) {
 	assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
 }
 
+func TestExtractIncidentClientFromHeaders(t *testing.T) {
+	t.Run("no headers, no env", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
+
+		client := IncidentClientFromContext(ctx)
+		require.NotNil(t, client)
+		assert.Equal(t, "http://localhost:3000/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
+	})
+
+	t.Run("no headers, with env", func(t *testing.T) {
+		t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
+
+		client := IncidentClientFromContext(ctx)
+		require.NotNil(t, client)
+		assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
+	})
+
+	t.Run("with headers, no env", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
+		require.NoError(t, err)
+		ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
+
+		client := IncidentClientFromContext(ctx)
+		require.NotNil(t, client)
+		assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
+	})
+
+	t.Run("with headers, with env", func(t *testing.T) {
+		t.Setenv("GRAFANA_URL", "will-not-be-used")
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
+		require.NoError(t, err)
+		ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
+
+		client := IncidentClientFromContext(ctx)
+		require.NotNil(t, client)
+		assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
+	})
+}
+
 func TestExtractGrafanaInfoFromHeaders(t *testing.T) {
 	t.Run("no headers, no env", func(t *testing.T) {
 		req, err := http.NewRequest("GET", "http://example.com", nil)
@@ -106,3 +153,64 @@ func TestExtractGrafanaClientPath(t *testing.T) {
 		assert.Equal(t, "/grafana/api", rt.BasePath)
 	})
 }
+
+// minURL is a helper struct representing what we can extract from a constructed
+// Grafana client.
+type minURL struct {
+	host, basePath string
+}
+
+// minURLFromClient extracts some minimal amount of URL info from a Grafana client.
+func minURLFromClient(c *grafana_client.GrafanaHTTPAPI) minURL {
+	rt := c.Transport.(*client.Runtime)
+	return minURL{rt.Host, rt.BasePath}
+}
+
+func TestExtractGrafanaClientFromHeaders(t *testing.T) {
+	t.Run("no headers, no env", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
+		c := GrafanaClientFromContext(ctx)
+		url := minURLFromClient(c)
+		assert.Equal(t, "localhost", url.host)
+		assert.Equal(t, "/api", url.basePath)
+	})
+
+	t.Run("no headers, with env", func(t *testing.T) {
+		t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
+
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
+		c := GrafanaClientFromContext(ctx)
+		url := minURLFromClient(c)
+		assert.Equal(t, "my-test-url.grafana.com", url.host)
+		assert.Equal(t, "/api", url.basePath)
+	})
+
+	t.Run("with headers, no env", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
+		ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
+		c := GrafanaClientFromContext(ctx)
+		url := minURLFromClient(c)
+		assert.Equal(t, "my-test-url.grafana.com", url.host)
+		assert.Equal(t, "/api", url.basePath)
+	})
+
+	t.Run("with headers, with env", func(t *testing.T) {
+		// Env vars should be ignored if headers are present.
+		t.Setenv("GRAFANA_URL", "will-not-be-used")
+
+		req, err := http.NewRequest("GET", "http://example.com", nil)
+		require.NoError(t, err)
+		req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
+		ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
+		c := GrafanaClientFromContext(ctx)
+		url := minURLFromClient(c)
+		assert.Equal(t, "my-test-url.grafana.com", url.host)
+		assert.Equal(t, "/api", url.basePath)
+	})
+}