Add GraphQLFeaturesTransport for reusable header handling

Co-authored-by: SamMorrowDrums <4811358+SamMorrowDrums@users.noreply.github.com>

Author: Copilot

internal/ghmcp/server.go

@@ -94,10 +94,13 @@ func createGitHubClients(cfg MCPServerConfig, apiHost apiHost) (*githubClients,
 
 	// Construct GraphQL client
 	// We use NewEnterpriseClient unconditionally since we already parsed the API host
+	// Layer transports: DefaultTransport -> bearerAuthTransport -> GraphQLFeaturesTransport
 	gqlHTTPClient := &http.Client{
-		Transport: &bearerAuthTransport{
-			transport: http.DefaultTransport,
-			token:     cfg.Token,
+		Transport: &github.GraphQLFeaturesTransport{
+			Transport: &bearerAuthTransport{
+				transport: http.DefaultTransport,
+				token:     cfg.Token,
+			},
 		},
 	}
 	gqlClient := githubv4.NewEnterpriseClient(apiHost.graphqlURL.String(), gqlHTTPClient)
@@ -622,12 +625,6 @@ type bearerAuthTransport struct {
 func (t *bearerAuthTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 	req = req.Clone(req.Context())
 	req.Header.Set("Authorization", "Bearer "+t.token)
-
-	// Check for GraphQL-Features in context and add header if present
-	if features := github.GetGraphQLFeatures(req.Context()); len(features) > 0 {
-		req.Header.Set("GraphQL-Features", strings.Join(features, ", "))
-	}
-
 	return t.transport.RoundTrip(req)
 }
 

pkg/github/transport.go

@@ -0,0 +1,44 @@
+package github
+
+import (
+	"net/http"
+	"strings"
+)
+
+// GraphQLFeaturesTransport is an http.RoundTripper that adds GraphQL-Features
+// header based on context values set by withGraphQLFeatures.
+//
+// This transport should be used in the HTTP client chain for githubv4.Client
+// to ensure GraphQL feature flags are properly sent to the GitHub API.
+//
+// Example usage:
+//
+//	httpClient := &http.Client{
+//		Transport: &github.GraphQLFeaturesTransport{
+//			Transport: http.DefaultTransport,
+//		},
+//	}
+//	gqlClient := githubv4.NewClient(httpClient)
+type GraphQLFeaturesTransport struct {
+	// Transport is the underlying http.RoundTripper. If nil, http.DefaultTransport is used.
+	Transport http.RoundTripper
+}
+
+// RoundTrip implements http.RoundTripper.
+// It adds the GraphQL-Features header if features are present in the request context.
+func (t *GraphQLFeaturesTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	transport := t.Transport
+	if transport == nil {
+		transport = http.DefaultTransport
+	}
+
+	// Clone request to avoid modifying the original
+	req = req.Clone(req.Context())
+
+	// Check for GraphQL-Features in context and add header if present
+	if features := GetGraphQLFeatures(req.Context()); len(features) > 0 {
+		req.Header.Set("GraphQL-Features", strings.Join(features, ", "))
+	}
+
+	return transport.RoundTrip(req)
+}

pkg/github/transport_test.go

@@ -0,0 +1,142 @@
+package github
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGraphQLFeaturesTransport(t *testing.T) {
+	tests := []struct {
+		name              string
+		features          []string
+		expectHeader      bool
+		expectedHeaderVal string
+	}{
+		{
+			name:              "adds single feature to header",
+			features:          []string{"issues_copilot_assignment_api_support"},
+			expectHeader:      true,
+			expectedHeaderVal: "issues_copilot_assignment_api_support",
+		},
+		{
+			name:              "adds multiple features to header",
+			features:          []string{"feature1", "feature2", "feature3"},
+			expectHeader:      true,
+			expectedHeaderVal: "feature1, feature2, feature3",
+		},
+		{
+			name:         "no header when no features in context",
+			features:     nil,
+			expectHeader: false,
+		},
+		{
+			name:         "no header when empty features slice",
+			features:     []string{},
+			expectHeader: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a test server that captures the request
+			var capturedReq *http.Request
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				capturedReq = r
+				w.WriteHeader(http.StatusOK)
+			}))
+			defer server.Close()
+
+			// Create HTTP client with GraphQLFeaturesTransport
+			client := &http.Client{
+				Transport: &GraphQLFeaturesTransport{
+					Transport: http.DefaultTransport,
+				},
+			}
+
+			// Create request with or without features in context
+			ctx := context.Background()
+			if tt.features != nil {
+				ctx = withGraphQLFeatures(ctx, tt.features...)
+			}
+
+			req, err := http.NewRequestWithContext(ctx, "GET", server.URL, nil)
+			require.NoError(t, err)
+
+			// Make request
+			resp, err := client.Do(req)
+			require.NoError(t, err)
+			defer resp.Body.Close()
+
+			// Verify header
+			if tt.expectHeader {
+				assert.Equal(t, tt.expectedHeaderVal, capturedReq.Header.Get("GraphQL-Features"))
+			} else {
+				assert.Empty(t, capturedReq.Header.Get("GraphQL-Features"))
+			}
+		})
+	}
+}
+
+func TestGraphQLFeaturesTransport_NilTransport(t *testing.T) {
+	// Test that nil Transport falls back to http.DefaultTransport
+	var capturedReq *http.Request
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedReq = r
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	client := &http.Client{
+		Transport: &GraphQLFeaturesTransport{
+			Transport: nil, // Explicitly nil
+		},
+	}
+
+	ctx := withGraphQLFeatures(context.Background(), "test_feature")
+	req, err := http.NewRequestWithContext(ctx, "GET", server.URL, nil)
+	require.NoError(t, err)
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	assert.Equal(t, "test_feature", capturedReq.Header.Get("GraphQL-Features"))
+}
+
+func TestGraphQLFeaturesTransport_PreservesOtherHeaders(t *testing.T) {
+	// Test that the transport doesn't interfere with other headers
+	var capturedReq *http.Request
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedReq = r
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	client := &http.Client{
+		Transport: &GraphQLFeaturesTransport{
+			Transport: http.DefaultTransport,
+		},
+	}
+
+	ctx := withGraphQLFeatures(context.Background(), "feature1")
+	req, err := http.NewRequestWithContext(ctx, "GET", server.URL, nil)
+	require.NoError(t, err)
+
+	// Add custom headers
+	req.Header.Set("Authorization", "Bearer test-token")
+	req.Header.Set("User-Agent", "test-agent")
+
+	resp, err := client.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	// Verify all headers are preserved
+	assert.Equal(t, "feature1", capturedReq.Header.Get("GraphQL-Features"))
+	assert.Equal(t, "Bearer test-token", capturedReq.Header.Get("Authorization"))
+	assert.Equal(t, "test-agent", capturedReq.Header.Get("User-Agent"))
+}