fix(router): ignore header rules for pubsub trigger sources

router-tests/events/trigger_test.go

@@ -1,12 +1,15 @@
 package events_test
 
 import (
+	"net/http"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 	"github.com/wundergraph/cosmo/router-tests/testenv"
+	"github.com/wundergraph/cosmo/router/core"
+	"github.com/wundergraph/cosmo/router/pkg/config"
 )
 
 // TestEDFSTriggerDeduplication verifies trigger ID generation for Cosmo Streams subscriptions.
@@ -33,8 +36,9 @@ func TestEDFSTriggerDeduplication(t *testing.T) {
 				// provider ID, so they must resolve to the same trigger ID — wait for exactly
 				// one trigger to be initialized before asserting.
 				xEnv.WaitForTriggerCount(1, time.Second*10)
-				xEnv.RequireTriggerCount(1)
 				xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename":"Employee"}`), 2, time.Second*10)
+				// Assert the exact count only after both subscriptions have been served
+				xEnv.RequireTriggerCount(1)
 			}()
 
 			// Subscription 1: selects only id.
@@ -104,29 +108,43 @@ func TestEDFSTriggerDeduplication(t *testing.T) {
 		})
 	})
 
-	// Two subscriptions with the same query but different initial_payload (headers) should
+	// Two subscriptions with the same query but different headers should
 	// still share a single NATS trigger because the trigger ID is based on the NATS subject,
 	// not on connection-level metadata like headers.
-	t.Run("same subject different initial payload shares one trigger", func(t *testing.T) {
+	t.Run("same subject different headers shares one trigger", func(t *testing.T) {
 		t.Parallel()
 		testenv.Run(t, &testenv.Config{
 			RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate,
 			EnableNats:               true,
+			RouterOptions: []core.Option{
+				core.WithHeaderRules(config.HeaderRules{
+					All: &config.GlobalHeaderRule{
+						Request: []*config.RequestHeaderRule{
+							{Operation: config.HeaderRuleOperationPropagate, Named: "Authorization"},
+						},
+					},
+				}),
+			},
 		}, func(t *testing.T, xEnv *testenv.Environment) {
 			var done sync.WaitGroup
 			done.Add(2)
 
 			go func() {
 				xEnv.WaitForSubscriptionCount(2, time.Second*10)
 				xEnv.WaitForTriggerCount(1, time.Second*10)
-				xEnv.RequireTriggerCount(1)
 				xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename":"Employee"}`), 2, time.Second*10)
+				// Asserted after both subscriptions have been served, see the comment in the
+				// "different selected fields" subtest above.
+				xEnv.RequireTriggerCount(1)
 			}()
 
-			// Subscription 1: sends Authorization header A in connection_init.
+			// Subscription 1: sends Authorization header A
 			go func() {
 				defer done.Done()
-				conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, []byte(`{"headers":{"Authorization":"Bearer token-a"}}`))
+				conn := xEnv.InitGraphQLWebSocketConnection(
+					http.Header{"Authorization": []string{"Bearer token-a"}}, nil,
+					[]byte(`{"headers":{"Authorization":"Bearer token-a"}}`),
+				)
 				defer conn.Close()
 
 				err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{
@@ -154,10 +172,13 @@ func TestEDFSTriggerDeduplication(t *testing.T) {
 				require.Equal(t, "1", complete.ID)
 			}()
 
-			// Subscription 2: sends a different Authorization header in connection_init.
+			// Subscription 2: sends Authorization header B
 			go func() {
 				defer done.Done()
-				conn := xEnv.InitGraphQLWebSocketConnection(nil, nil, []byte(`{"headers":{"Authorization":"Bearer token-b"}}`))
+				conn := xEnv.InitGraphQLWebSocketConnection(
+					http.Header{"Authorization": []string{"Bearer token-b"}}, nil,
+					[]byte(`{"headers":{"Authorization":"Bearer token-b"}}`),
+				)
 				defer conn.Close()
 
 				err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{

router-tests/operations/singleflight_test.go

@@ -1,7 +1,6 @@
 package integration
 
 import (
-	"context"
 	"fmt"
 	"net/http"
 	"sync"
@@ -521,7 +520,11 @@ func TestSingleFlight(t *testing.T) {
 			require.Less(t, actualSubgraphRequests, numOfOperations)
 		})
 	})
-	t.Run("subscription deduplication with multiple subgraphs - different headers", func(t *testing.T) {
+	// All subscriptions share a single EDFS trigger regardless of their headers, because the
+	// trigger ID of a pubsub source is derived from the subject and provider only. The nested
+	// fetches each subscription performs to resolve its response are still built from the
+	// per-request propagated headers, so those must not be de-duplicated by single flight.
+	t.Run("subscription with different headers does not deduplicate subgraph fetches", func(t *testing.T) {
 		t.Parallel()
 		testenv.Run(t, &testenv.Config{
 			RouterConfigJSONTemplate: testenv.ConfigWithEdfsNatsJSONTemplate,
@@ -556,24 +559,14 @@ func TestSingleFlight(t *testing.T) {
 			)
 			done.Add(int(numOfOperations))
 
-			// Continuously publish until all consumers have received their message.
-			// NATSPublishUntilMinMessagesSent is insufficient here because cumulative
-			// MessagesSent can reach 10 before all 10 consumers are served (retries
-			// deliver to already-served consumers, inflating the count).
-			publishCtx, publishCancel := context.WithCancel(xEnv.Context)
+			// Wait for all subscriptions to be established before triggering. The differing
+			// Authorization headers do not split the trigger, so a single message fans out
+			// to all subscriptions.
 			go func() {
 				xEnv.WaitForSubscriptionCount(uint64(numOfOperations), time.Second*15)
-				xEnv.WaitForTriggerCount(uint64(numOfOperations), time.Second*15)
-				for {
-					select {
-					case <-publishCtx.Done():
-						return
-					default:
-					}
-					_ = xEnv.NatsConnectionDefault.Publish(xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename": "Employee"}`))
-					_ = xEnv.NatsConnectionDefault.Flush()
-					time.Sleep(500 * time.Millisecond)
-				}
+				xEnv.WaitForTriggerCount(1, time.Second*15)
+				// Trigger the subscription via NATS to get updates for all subscriptions
+				xEnv.NATSPublishUntilReceived(xEnv.NatsConnectionDefault, xEnv.GetPubSubName("employeeUpdated.3"), []byte(`{"id":3,"__typename": "Employee"}`), 1, time.Second*15)
 			}()
 
 			for i := int64(0); i < numOfOperations; i++ {
@@ -606,29 +599,26 @@ func TestSingleFlight(t *testing.T) {
 					})
 					require.NoError(t, err)
 
-					// Read messages until we get "complete", draining any extra
-					// "next" messages that may arrive from publish retries
-					for {
-						var reply testenv.WebSocketMessage
-						err = conn.SetReadDeadline(time.Now().Add(2 * time.Second))
-						require.NoError(t, err)
-						err = testenv.WSReadJSON(t, conn, &reply)
-						require.NoError(t, err)
-						if reply.Type == "complete" {
-							require.Equal(t, "1", reply.ID)
-							break
-						}
-					}
+					// Read the complete message
+					var complete testenv.WebSocketMessage
+					err = conn.SetReadDeadline(time.Now().Add(1 * time.Second))
+					require.NoError(t, err)
+					err = testenv.WSReadJSON(t, conn, &complete)
+					require.NoError(t, err)
+					require.Equal(t, "complete", complete.Type)
+					require.Equal(t, "1", complete.ID)
 				}(i)
 			}
 			done.Wait()
-			publishCancel()
 			xEnv.WaitForSubscriptionCount(0, time.Second*5)
 
-			// We expect no request de-duplication because different headers must not be de-duplicated
-			// Publish retries may increase the count, so check >= not ==
+			// We expect no request de-duplication because the fetches carry different headers.
+			// The NATS event itself supplies __typename and id — the only fields the pubsub
+			// data source owns — so resolving details.forename and details.surname costs one
+			// entity fetch to the employees subgraph per subscription: 10 in total.
 			actualSubgraphRequests := xEnv.SubgraphRequestCount.Global.Load()
-			require.GreaterOrEqual(t, actualSubgraphRequests, numOfOperations)
+			require.Equal(t, numOfOperations, actualSubgraphRequests)
+			require.Equal(t, numOfOperations, xEnv.SubgraphRequestCount.Employees.Load())
 		})
 	})
 	t.Run("mutation with multiple subgraphs deduplication", func(t *testing.T) {

router-tests/subscriptions/websocket_test.go

@@ -741,6 +741,90 @@ func TestWebSockets(t *testing.T) {
 			xEnv.WaitForSubscriptionCount(0, time.Second*5)
 		})
 	})
+	// Trigger IDs for regular (non-EDFS) subgraph subscriptions include the hash of the
+	// headers built by the SubgraphHeadersBuilder, so clients whose propagated headers
+	// differ must never share an upstream trigger. Upgrade header, query param and initial
+	// payload forwarding are disabled in both subtests so that the trigger input is identical
+	// for both clients and header propagation is the only discriminator.
+	t.Run("subscription trigger deduplication with header propagation", func(t *testing.T) {
+		t.Parallel()
+
+		headerRules := config.HeaderRules{
+			All: &config.GlobalHeaderRule{
+				Request: []*config.RequestHeaderRule{
+					{
+						Operation: config.HeaderRuleOperationPropagate,
+						Named:     "Authorization",
+					},
+				},
+			},
+		}
+
+		websocketConfig := func(cfg *config.WebSocketConfiguration) {
+			cfg.ForwardUpgradeHeaders.Enabled = false
+			cfg.ForwardUpgradeQueryParams.Enabled = false
+			cfg.ForwardInitialPayload = false
+		}
+
+		subscribeCurrentTime := func(t *testing.T, xEnv *testenv.Environment, authorization string) *websocket.Conn {
+			t.Helper()
+
+			conn := xEnv.InitGraphQLWebSocketConnection(http.Header{
+				"Authorization": []string{authorization},
+			}, nil, nil)
+
+			err := testenv.WSWriteJSON(t, conn, &testenv.WebSocketMessage{
+				ID:      "1",
+				Type:    "subscribe",
+				Payload: []byte(`{"query":"subscription { currentTime { unixTime timeStamp }}"}`),
+			})
+			require.NoError(t, err)
+
+			return conn
+		}
+
+		t.Run("different headers use separate triggers", func(t *testing.T) {
+			t.Parallel()
+
+			testenv.Run(t, &testenv.Config{
+				ModifyWebsocketConfiguration: websocketConfig,
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(headerRules),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				connA := subscribeCurrentTime(t, xEnv, "Bearer token-a")
+				defer connA.Close()
+
+				connB := subscribeCurrentTime(t, xEnv, "Bearer token-b")
+				defer connB.Close()
+
+				xEnv.WaitForSubscriptionCount(2, time.Second*15)
+				xEnv.WaitForTriggerCount(2, time.Second*15)
+				xEnv.RequireTriggerCount(2)
+			})
+		})
+
+		t.Run("same headers share one trigger", func(t *testing.T) {
+			t.Parallel()
+
+			testenv.Run(t, &testenv.Config{
+				ModifyWebsocketConfiguration: websocketConfig,
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(headerRules),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				connA := subscribeCurrentTime(t, xEnv, "Bearer token-a")
+				defer connA.Close()
+
+				connB := subscribeCurrentTime(t, xEnv, "Bearer token-a")
+				defer connB.Close()
+
+				xEnv.WaitForSubscriptionCount(2, time.Second*15)
+				xEnv.WaitForTriggerCount(1, time.Second*15)
+				xEnv.RequireTriggerCount(1)
+			})
+		})
+	})
 	t.Run("empty allow lists should allow all headers and query args", func(t *testing.T) {
 		t.Parallel()
 

router/core/context.go

@@ -24,6 +24,7 @@ import (
 	"github.com/wundergraph/cosmo/router/pkg/authentication"
 	"github.com/wundergraph/cosmo/router/pkg/config"
 	"github.com/wundergraph/cosmo/router/pkg/graphqlschemausage"
+	pubsub "github.com/wundergraph/cosmo/router/pkg/pubsub/datasource"
 	ctrace "github.com/wundergraph/cosmo/router/pkg/trace"
 
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/datasource/httpclient"
@@ -334,6 +335,15 @@ func SubgraphHeadersBuilder(ctx *requestContext, headerPropagation *HeaderPropag
 		headers := make(map[string]*HeaderWithHash, len(p.Response.Response.DataSources)+1)
 		makeHeaders(headers, p.Response.Response.DataSources)
 
+		// avoid adding header rules for pubsub triggers sources, as no headers are passed to them.
+		_, isPubSub := p.Response.Trigger.Source.(pubsub.SubscriptionDataSource)
+		if isPubSub {
+			return &headerBuilder{
+				headers: headers,
+				allHash: keyGen.Sum64(),
+			}
+		}
+
 		h, hh := headerPropagation.BuildRequestHeaderForSubgraph(p.Response.Trigger.SourceName, ctx)
 		headers[p.Response.Trigger.SourceName] = &HeaderWithHash{
 			Header: h,