feat(router): allow split config polling

router/core/init_config_poller.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 
+	rjwt "github.com/wundergraph/cosmo/router/internal/jwt"
 	"github.com/wundergraph/cosmo/router/pkg/controlplane/configpoller"
 	"github.com/wundergraph/cosmo/router/pkg/execution_config"
 	"github.com/wundergraph/cosmo/router/pkg/routerconfig"
@@ -125,6 +126,22 @@ func InitializeConfigPoller(r *Router, registry *ProviderRegistry) (*configpolle
 		return nil, nil
 	}
 
+	// Check whether the router JWT requests the split-config-loading strategy.
+	// Split config is only supported with the default Cosmo CDN (no custom storage provider).
+	hasSplitCfgFeature, err := hasSplitConfigFeature(r)
+	if err != nil {
+		return nil, err
+	}
+
+	if hasSplitCfgFeature {
+		providerID := r.routerConfigPollerConfig.Storage.ProviderID
+		if providerID == "" {
+			return newSplitConfigPoller(r)
+		}
+		r.logger.Info("split-config-loading feature is enabled but a custom storage provider is configured; falling back to regular config polling",
+			zap.String("provider_id", providerID))
+	}
+
 	primaryClient, err := getConfigClient(r, registry, r.routerConfigPollerConfig.Storage.ProviderID, false)
 	if err != nil {
 		return nil, err
@@ -159,3 +176,37 @@ func InitializeConfigPoller(r *Router, registry *ProviderRegistry) (*configpolle
 
 	return &configPoller, nil
 }
+
+func hasSplitConfigFeature(r *Router) (bool, error) {
+	if r.graphApiToken == "" {
+		return false, nil
+	}
+
+	claims, err := rjwt.ExtractFederatedGraphTokenClaims(r.graphApiToken)
+	if err != nil {
+		return false, fmt.Errorf("failed to parse graph API token: %w", err)
+	}
+
+	return claims.HasFeature(rjwt.FeatureSplitConfigLoading), nil
+}
+
+func newSplitConfigPoller(r *Router) (*configpoller.ConfigPoller, error) {
+	fetcher, err := configCDNProvider.NewSplitFetcher(
+		r.cdnConfig.URL,
+		r.graphApiToken,
+		&configCDNProvider.Options{
+			Logger:       r.logger,
+			SignatureKey: r.routerConfigPollerConfig.GraphSignKey,
+		},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create split config fetcher: %w", err)
+	}
+
+	splitPoller := configpoller.NewSplitConfigPoller(
+		fetcher,
+		configpoller.WithSplitLogger(r.logger),
+		configpoller.WithSplitPolling(r.routerConfigPollerConfig.PollInterval, r.routerConfigPollerConfig.PollJitter),
+	)
+	return &splitPoller, nil
+}

router/internal/jwt/claims.go

@@ -2,17 +2,27 @@ package jwt
 
 import (
 	"fmt"
+	"slices"
+
 	"github.com/golang-jwt/jwt/v5"
 )
 
 const (
-	FederatedGraphIDClaim = "federated_graph_id"
-	OrganizationIDClaim   = "organization_id"
+	FederatedGraphIDClaim     = "federated_graph_id"
+	OrganizationIDClaim       = "organization_id"
+	FeaturesClaim             = "features"
+	FeatureSplitConfigLoading = "split-config-loading"
 )
 
 type FederatedGraphTokenClaims struct {
 	FederatedGraphID string
 	OrganizationID   string
+	Features         []string
+}
+
+// HasFeature reports whether the given feature flag is present in the token claims.
+func (c *FederatedGraphTokenClaims) HasFeature(feature string) bool {
+	return slices.Contains(c.Features, feature)
 }
 
 func ExtractFederatedGraphTokenClaims(token string) (*FederatedGraphTokenClaims, error) {
@@ -44,8 +54,24 @@ func ExtractFederatedGraphTokenClaims(token string) (*FederatedGraphTokenClaims,
 		return nil, fmt.Errorf("invalid token claims, %q is not a string, it's %T", OrganizationIDClaim, organizationIDValue)
 	}
 
+	var features []string
+	if featuresValue := claims[FeaturesClaim]; featuresValue != nil {
+		featuresSlice, ok := featuresValue.([]any)
+		if !ok {
+			return nil, fmt.Errorf("invalid token claims, %q is not an array, it's %T", FeaturesClaim, featuresValue)
+		}
+		for _, v := range featuresSlice {
+			s, ok := v.(string)
+			if !ok {
+				return nil, fmt.Errorf("invalid token claims, %q contains a non-string element %T", FeaturesClaim, v)
+			}
+			features = append(features, s)
+		}
+	}
+
 	return &FederatedGraphTokenClaims{
 		FederatedGraphID: federatedGraphID,
 		OrganizationID:   organizationID,
+		Features:         features,
 	}, nil
 }