feat: use cors configuration also on mcp server

router-tests/mcp_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/wundergraph/cosmo/router-tests/testenv"
 	"github.com/wundergraph/cosmo/router/core"
 	"github.com/wundergraph/cosmo/router/pkg/config"
+	"github.com/wundergraph/cosmo/router/pkg/cors"
 	"github.com/wundergraph/cosmo/router/pkg/schemaloader"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/astparser"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/asttransform"
@@ -559,6 +560,56 @@ func TestMCP(t *testing.T) {
 				}
 			})
 		})
+
+		t.Run("Custom headers in router CORS are also used by the MCP server", func(t *testing.T) {
+			testenv.Run(t, &testenv.Config{
+				MCP: config.MCPConfiguration{
+					Enabled: true,
+				},
+				RouterOptions: []core.Option{
+					core.WithCors(&cors.Config{
+						AllowHeaders: []string{"Test", "X-Custom-Auth"},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				// Get the MCP server address from the configuration
+				mcpAddr := xEnv.GetMCPServerAddr()
+
+				// Create a GET request
+				req, err := http.NewRequest("GET", mcpAddr, nil)
+				require.NoError(t, err)
+
+				// Add cross-origin header
+				req.Header.Set("Origin", "https://example.com")
+
+				// Make the request
+				resp, err := xEnv.RouterClient.Do(req)
+				require.NoError(t, err)
+				defer resp.Body.Close()
+
+				// Verify CORS headers are present in the response
+				assert.Equal(t, "*", resp.Header.Get("Access-Control-Allow-Origin"))
+
+				allowedMethods := resp.Header.Get("Access-Control-Allow-Methods")
+				assert.Contains(t, allowedMethods, "GET")
+				assert.Contains(t, allowedMethods, "POST")
+				assert.Contains(t, allowedMethods, "PUT")
+				assert.Contains(t, allowedMethods, "DELETE")
+				assert.Contains(t, allowedMethods, "OPTIONS")
+
+				allowedHeaders := resp.Header.Get("Access-Control-Allow-Headers")
+				assert.Contains(t, allowedHeaders, "Content-Type")
+				assert.Contains(t, allowedHeaders, "Accept")
+				assert.Contains(t, allowedHeaders, "Authorization")
+				assert.Contains(t, allowedHeaders, "Last-Event-ID")
+				assert.Contains(t, allowedHeaders, "Mcp-Protocol-Version")
+				assert.Contains(t, allowedHeaders, "Mcp-Session-Id")
+				assert.Contains(t, allowedHeaders, "Test")
+				assert.Contains(t, allowedHeaders, "X-Custom-Auth")
+
+				assert.Equal(t, "86400", resp.Header.Get("Access-Control-Max-Age"))
+			})
+		})
 	})
 
 	t.Run("Operation Description Extraction", func(t *testing.T) {

router/core/router.go

@@ -916,6 +916,10 @@ func (r *Router) bootstrap(ctx context.Context) error {
 			mcpserver.WithStateless(r.mcp.Session.Stateless),
 		}
 
+		if r.corsOptions != nil {
+			mcpOpts = append(mcpOpts, mcpserver.WithCORS(*r.corsOptions))
+		}
+
 		// Determine the router GraphQL endpoint
 		var routerGraphQLEndpoint string
 

router/pkg/mcpserver/server.go

@@ -16,6 +16,8 @@ import (
 	"github.com/mark3labs/mcp-go/mcp"
 	"github.com/mark3labs/mcp-go/server"
 	"github.com/santhosh-tekuri/jsonschema/v6"
+	"github.com/wundergraph/cosmo/router/internal/stringsx"
+	"github.com/wundergraph/cosmo/router/pkg/cors"
 	"github.com/wundergraph/cosmo/router/pkg/schemaloader"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/ast"
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/astprinter"
@@ -91,6 +93,8 @@ type Options struct {
 	ExposeSchema bool
 	// Stateless determines whether the MCP server should be stateless
 	Stateless bool
+	// CorsConfig is the CORS configuration for the MCP server
+	CorsConfig cors.Config
 }
 
 // GraphQLSchemaServer represents an MCP server that works with GraphQL schemas and operations
@@ -111,6 +115,7 @@ type GraphQLSchemaServer struct {
 	operationsManager         *OperationsManager
 	schemaCompiler            *SchemaCompiler
 	registeredTools           []string
+	corsConfig                cors.Config
 }
 
 type graphqlRequest struct {
@@ -237,6 +242,7 @@ func NewGraphQLSchemaServer(routerGraphQLEndpoint string, opts ...func(*Options)
 		enableArbitraryOperations: options.EnableArbitraryOperations,
 		exposeSchema:              options.ExposeSchema,
 		stateless:                 options.Stateless,
+		corsConfig:                options.CorsConfig,
 	}
 
 	return gs, nil
@@ -302,6 +308,12 @@ func WithStateless(stateless bool) func(*Options) {
 	}
 }
 
+func WithCORS(corsCfg cors.Config) func(*Options) {
+	return func(o *Options) {
+		o.CorsConfig = corsCfg
+	}
+}
+
 // Serve starts the server with the configured options and returns a streamable HTTP server.
 func (s *GraphQLSchemaServer) Serve() (*server.StreamableHTTPServer, error) {
 	// Create custom HTTP server
@@ -320,12 +332,12 @@ func (s *GraphQLSchemaServer) Serve() (*server.StreamableHTTPServer, error) {
 		server.WithHeartbeatInterval(10*time.Second),
 	)
 
-	corsMiddleware := WithCORS("GET", "POST", "PUT", "DELETE")
+	middleware := corsMiddleware(s.corsConfig)
 
 	mux := http.NewServeMux()
 
 	// No OAuth protection - original behavior
-	mux.Handle("/mcp", corsMiddleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	mux.Handle("/mcp", middleware(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		streamableHTTPServer.ServeHTTP(w, r)
 	})))
 
@@ -789,24 +801,24 @@ func (s *GraphQLSchemaServer) handleGetGraphQLSchema() func(ctx context.Context,
 	}
 }
 
-// WithCORS creates a reusable CORS middleware that can be used with any HTTP handler.
+// corsMiddleware creates a reusable CORS middleware that can be used with any HTTP handler.
 // It handles preflight OPTIONS requests and sets appropriate CORS headers.
 //
 // Example usage:
 //
-//	corsMiddleware := WithCORS("GET", "POST", "PUT", "DELETE")
-//	http.Handle("/api/", corsMiddleware(apiHandler))
+//	middleware := corsMiddleware("GET", "POST", "PUT", "DELETE")
+//	http.Handle("/api/", middleware(apiHandler))
 //
 // The middleware sets the following CORS headers:
 //   - Access-Control-Allow-Origin: *
 //   - Access-Control-Allow-Methods: specified methods + OPTIONS
 //   - Access-Control-Allow-Headers: Content-Type, Authorization
 //   - Access-Control-Max-Age: 86400 (24 hours)
-func WithCORS(allowedMethods ...string) func(http.Handler) http.Handler {
+func corsMiddleware(corsConfig cors.Config) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 			// Set CORS headers for all requests
-			setCORSHeaders(w, allowedMethods)
+			setCORSHeaders(w, corsConfig)
 
 			// Handle preflight OPTIONS requests
 			if req.Method == http.MethodOptions {
@@ -822,9 +834,14 @@ func WithCORS(allowedMethods ...string) func(http.Handler) http.Handler {
 
 // setCORSHeaders sets common CORS headers
 // Only used for web browsers, not for API clients
-func setCORSHeaders(w http.ResponseWriter, allowedMethods []string) {
+func setCORSHeaders(w http.ResponseWriter, config cors.Config) {
 	w.Header().Set("Access-Control-Allow-Origin", "*")
-	w.Header().Set("Access-Control-Allow-Methods", strings.Join(append(allowedMethods, "OPTIONS"), ", "))
-	w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Accept, Authorization, Last-Event-ID, Mcp-Protocol-Version, Mcp-Session-Id")
-	w.Header().Set("Access-Control-Max-Age", "86400") // 24 hours
+	w.Header().Set("Access-Control-Allow-Methods", "GET, PUT, POST, DELETE, OPTIONS")
+	hdrs := stringsx.RemoveDuplicates(append(config.AllowHeaders, "Content-Type", "Accept", "Authorization", "Last-Event-ID", "Mcp-Protocol-Version", "Mcp-Session-Id"))
+	w.Header().Set("Access-Control-Allow-Headers", strings.Join(hdrs, ", "))
+	maxAge := config.MaxAge
+	if maxAge <= 0 {
+		maxAge = 24 * time.Hour
+	}
+	w.Header().Set("Access-Control-Max-Age", fmt.Sprintf("%d", int(maxAge.Seconds())))
 }