@@ -934,17 +934,6 @@ function federateSubgraphs(supergraph: Schema, subgraphs: QueryGraph[]): QueryGr
934934 }
935935 }
936936
937- for ( const [ subgraphName , args ] of subgraphToArgs ) {
938- args . sort ( ) ;
939- const argToIndex = new Map ( ) ;
940- for ( let idx = 0 ; idx < args . length ; idx ++ ) {
941- argToIndex . set ( args [ idx ] , `contextualArgument_${ i } _${ idx } ` ) ;
942- }
943- subgraphToArgIndices . set ( subgraphName , argToIndex ) ;
944- }
945-
946- builder . setContextMaps ( subgraphToArgs , subgraphToArgIndices ) ;
947-
948937 simpleTraversal (
949938 subgraph ,
950939 _v => { return undefined ; } ,
@@ -965,6 +954,22 @@ function federateSubgraphs(supergraph: Schema, subgraphs: QueryGraph[]): QueryGr
965954
966955 }
967956
957+ // add contextual argument maps to builder
958+ for ( const [ i , subgraph ] of subgraphs . entries ( ) ) {
959+ const subgraphName = subgraph . name ;
960+ const args = subgraphToArgs . get ( subgraph . name ) ;
961+ if ( args ) {
962+ args . sort ( ) ;
963+ const argToIndex = new Map ( ) ;
964+ for ( let idx = 0 ; idx < args . length ; idx ++ ) {
965+ argToIndex . set ( args [ idx ] , `contextualArgument_${ i + 1 } _${ idx } ` ) ;
966+ }
967+ subgraphToArgIndices . set ( subgraphName , argToIndex ) ;
968+ }
969+ }
970+
971+ builder . setContextMaps ( subgraphToArgs , subgraphToArgIndices ) ;
972+
968973 // Now we handle @provides
969974 let provideId = 0 ;
970975 for ( const [ i , subgraph ] of subgraphs . entries ( ) ) {
0 commit comments