feat(router): support costs on arguments of directives

composition/src/router-configuration/types.ts

@@ -109,6 +109,7 @@ export type Costs = {
 
 export type FieldWeightConfiguration = {
   argumentWeights: Map<ArgumentName, number>;
+  directiveArgumentWeights: Map<DirectiveArgumentCoords, number>;
   fieldName: FieldName;
   typeName: TypeName;
   weight?: number;

composition/src/v1/normalization/normalization-factory.ts

@@ -212,6 +212,7 @@ import {
   type ConfigureDescriptionData,
   type EntityData,
   type EntityInterfaceSubgraphData,
+  type DirectiveDefinitionData,
   type EnumDefinitionData,
   type EnumValueData,
   ExtensionType,
@@ -378,6 +379,7 @@ import {
   type FieldSetParentResult,
   type HandleCostDirectiveParams,
   type HandleListSizeDirectiveParams,
+  type RecordDirectiveWeightOnFieldParams,
   type HandleOverrideDirectiveParams,
   type HandleRequiresScopesDirectiveParams,
   type HandleSemanticNonNullDirectiveParams,
@@ -750,6 +752,9 @@ export class NormalizationFactory {
           listSizeFieldMustReturnListOrUseSizedFieldsErrorMessage(directiveCoords, printTypeNode(data.type)),
         );
       }
+      if (!isCost && !isListSize && isField) {
+        this.recordDirectiveWeightOnField({ data: data as FieldData, definitionData, directiveName, directiveNode });
+      }
       return errorMessages;
     }
     const definedArgumentNames = new Set<string>();
@@ -815,6 +820,9 @@ export class NormalizationFactory {
     } else if (isListSize && isField) {
       this.handleListSizeDirective({ data, directiveCoords, directiveNode, errorMessages });
     }
+    if (!isCost && !isListSize && isField) {
+      this.recordDirectiveWeightOnField({ data: data as FieldData, definitionData, directiveName, directiveNode });
+    }
     if (duplicateArgumentNames.size > 0) {
       errorMessages.push(duplicateDirectiveArgumentDefinitionsErrorMessage([...duplicateArgumentNames]));
     }
@@ -2472,6 +2480,20 @@ export class NormalizationFactory {
     data.nullLevelsBySubgraphName.set(this.subgraphName, levels);
   }
 
+  getOrCreateFieldWeight(typeName: TypeName, fieldName: FieldName): FieldWeightConfiguration {
+    const fieldCoords = `${typeName}.${fieldName}`;
+    return getValueOrDefault(
+      this.costs.fieldWeights,
+      fieldCoords,
+      (): FieldWeightConfiguration => ({
+        typeName,
+        fieldName,
+        argumentWeights: new Map(),
+        directiveArgumentWeights: new Map(),
+      }),
+    );
+  }
+
   handleCostDirective({ data, directiveCoords, directiveNode, errorMessages }: HandleCostDirectiveParams) {
     const weightArg = directiveNode.arguments?.find((arg) => arg.name.value === WEIGHT);
     if (!weightArg || weightArg.value.kind !== Kind.INT) {
@@ -2495,16 +2517,7 @@ export class NormalizationFactory {
           errorMessages.push(costOnInterfaceFieldErrorMessage(directiveCoords));
           break;
         }
-        const fieldCoords = `${typeName}.${data.name}`;
-        const fieldWeight = getValueOrDefault(
-          this.costs.fieldWeights,
-          fieldCoords,
-          (): FieldWeightConfiguration => ({
-            typeName,
-            fieldName: data.name,
-            argumentWeights: new Map(),
-          }),
-        );
+        const fieldWeight = this.getOrCreateFieldWeight(typeName, data.name);
         fieldWeight.weight = weightValue;
         break;
       }
@@ -2522,36 +2535,61 @@ export class NormalizationFactory {
             errorMessages.push(costOnInterfaceFieldErrorMessage(directiveCoords));
             break;
           }
-          const parentFieldCoords = `${typeName}.${ivData.fieldName}`;
-          const fieldWeight = getValueOrDefault(
-            this.costs.fieldWeights,
-            parentFieldCoords,
-            (): FieldWeightConfiguration => ({
-              typeName,
-              fieldName: ivData.fieldName!,
-              argumentWeights: new Map(),
-            }),
-          );
+          const fieldWeight = this.getOrCreateFieldWeight(typeName, ivData.fieldName!);
           fieldWeight.argumentWeights.set(ivData.name, weightValue);
         } else {
           const typeName = ivData.renamedParentTypeName || ivData.originalParentTypeName;
-          const fieldCoords = `${typeName}.${ivData.name}`;
-          const fieldWeight = getValueOrDefault(
-            this.costs.fieldWeights,
-            fieldCoords,
-            (): FieldWeightConfiguration => ({
-              typeName,
-              fieldName: ivData.name,
-              argumentWeights: new Map(),
-            }),
-          );
+          const fieldWeight = this.getOrCreateFieldWeight(typeName, ivData.name);
           fieldWeight.weight = weightValue;
         }
         break;
       }
     }
   }
 
+  recordDirectiveWeightOnField({
+    data,
+    definitionData,
+    directiveName,
+    directiveNode,
+  }: RecordDirectiveWeightOnFieldParams) {
+    const typeName = data.renamedParentTypeName || data.originalParentTypeName;
+    const parentTypeData = this.parentDefinitionDataByTypeName.get(typeName);
+    // Directive argument weights should only be on concrete type fields.
+    if (!parentTypeData || parentTypeData.kind === Kind.INTERFACE_TYPE_DEFINITION) {
+      return;
+    }
+
+    // Determine which arguments are active (non-null) on this directive usage.
+    // An argument is active if it has an explicit non-null value or
+    // if it has a default value and was not explicitly set to null.
+    const explicitArgs = new Map<string, ConstValueNode>();
+    if (directiveNode.arguments) {
+      for (const arg of directiveNode.arguments) {
+        explicitArgs.set(arg.name.value, arg.value);
+      }
+    }
+
+    for (const [argName, argData] of definitionData.argumentTypeNodeByName) {
+      const coords = `${directiveName}.${argName}`;
+      const argWeight = this.costs.directiveArgumentWeights.get(coords);
+      if (argWeight === undefined) {
+        continue;
+      }
+      // Check if this argument is active (non-null) at the usage site
+      const explicitValue = explicitArgs.get(argName);
+      if (explicitValue) {
+        if (explicitValue.kind === Kind.NULL) {
+          continue;
+        }
+      } else if (!argData.defaultValue || argData.defaultValue.kind === Kind.NULL) {
+        continue;
+      }
+      const fieldWeight = this.getOrCreateFieldWeight(typeName, data.name);
+      fieldWeight.directiveArgumentWeights.set(coords, argWeight);
+    }
+  }
+
   handleListSizeDirective({ data, directiveCoords, directiveNode, errorMessages }: HandleListSizeDirectiveParams) {
     const args = directiveNode.arguments;
     if (!args) {

composition/src/v1/normalization/types.ts

@@ -82,6 +82,13 @@ export type HandleListSizeDirectiveParams = {
   errorMessages: Array<string>;
 };
 
+export type RecordDirectiveWeightOnFieldParams = {
+  data: FieldData;
+  definitionData: DirectiveDefinitionData;
+  directiveName: string;
+  directiveNode: ConstDirectiveNode;
+};
+
 export type AddInputValueDataByNodeParams = {
   inputValueDataByName: Map<string, InputValueData>;
   isArgument: boolean;

composition/tests/v1/directives/cost.test.ts

@@ -2,13 +2,15 @@ import { describe, expect, test } from 'vitest';
 import {
   COST,
   costOnInterfaceFieldErrorMessage,
+  type FieldName,
   FIRST_ORDINAL,
   invalidArgumentValueErrorMessage,
   invalidDirectiveError,
   invalidDirectiveLocationErrorMessage,
   parse,
   ROUTER_COMPATIBILITY_VERSION_ONE,
   type Subgraph,
+  type TypeName,
   undefinedRequiredArgumentsErrorMessage,
 } from '../../../src';
 import { COST_DIRECTIVE, SCHEMA_QUERY_DEFINITION } from '../utils/utils';
@@ -20,6 +22,16 @@ import {
   schemaToSortedNormalizedString,
 } from '../../utils/utils';
 
+function fieldWeight(typeName: TypeName, fieldName: FieldName, overrides: Record<string, unknown> = {}) {
+  return {
+    typeName,
+    fieldName,
+    argumentWeights: new Map(),
+    directiveArgumentWeights: new Map(),
+    ...overrides,
+  };
+}
+
 const NORMALIZATION_SCHEMA_QUERY = `
   schema {
     query: Query
@@ -367,18 +379,8 @@ describe('@cost directive tests', () => {
         const costsA = subgraphConfigBySubgraphName.get('subgraph-cost-shared-a')?.costs;
         const costsB = subgraphConfigBySubgraphName.get('subgraph-cost-shared-b')?.costs;
 
-        expect(costsA?.fieldWeights.get('User.name')).toEqual({
-          typeName: 'User',
-          fieldName: 'name',
-          argumentWeights: new Map(),
-          weight: 10,
-        });
-        expect(costsB?.fieldWeights.get('User.name')).toEqual({
-          typeName: 'User',
-          fieldName: 'name',
-          argumentWeights: new Map(),
-          weight: 20,
-        });
+        expect(costsA?.fieldWeights.get('User.name')).toEqual(fieldWeight('User', 'name', { weight: 10 }));
+        expect(costsB?.fieldWeights.get('User.name')).toEqual(fieldWeight('User', 'name', { weight: 20 }));
       });
 
       test('that fields unique to each subgraph only appear in their respective cost configs', () => {
@@ -491,12 +493,7 @@ describe('@cost directive tests', () => {
         subgraphWithCostOnImplementingTypeField,
         ROUTER_COMPATIBILITY_VERSION_ONE,
       );
-      expect(costs.fieldWeights.get('User.id')).toEqual({
-        typeName: 'User',
-        fieldName: 'id',
-        argumentWeights: new Map(),
-        weight: 5,
-      });
+      expect(costs.fieldWeights.get('User.id')).toEqual(fieldWeight('User', 'id', { weight: 5 }));
     });
   });
 
@@ -517,6 +514,66 @@ describe('@cost directive tests', () => {
       expect(costs.directiveArgumentWeights.get('myDirective.arg2')).toBe(7);
     });
 
+    test('that a field with a cost-weighted directive records directiveArgumentWeights', () => {
+      const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnDirectiveArgument, ROUTER_COMPATIBILITY_VERSION_ONE);
+      // Query.field has @myDirective(arg1: "hello"), and arg1 has @cost(weight: 5)
+      expect(costs.fieldWeights.get('Query.field')).toEqual(
+        fieldWeight('Query', 'field', { directiveArgumentWeights: new Map([['myDirective.arg1', 5]]) }),
+      );
+    });
+
+    test('that a field with a multi-arg cost-weighted directive stores per-argument weights', () => {
+      const { costs } = normalizeSubgraphSuccess(
+        subgraphWithCostOnMultipleDirectiveArguments,
+        ROUTER_COMPATIBILITY_VERSION_ONE,
+      );
+      // Query.field has @myDirective(arg1: "hello", arg2: 42), arg1 costs 3, arg2 costs 7
+      expect(costs.fieldWeights.get('Query.field')).toEqual(
+        fieldWeight('Query', 'field', {
+          directiveArgumentWeights: new Map([
+            ['myDirective.arg1', 3],
+            ['myDirective.arg2', 7],
+          ]),
+        }),
+      );
+    });
+
+    test('that a field with a directive whose argument is null does not record directiveArgumentWeights', () => {
+      const { costs } = normalizeSubgraphSuccess(
+        {
+          name: 'subgraph-null-directive-arg',
+          url: '',
+          definitions: parse(`
+          directive @myDirective(arg1: String @cost(weight: 5)) on FIELD_DEFINITION
+          type Query {
+            field: String! @myDirective(arg1: null)
+          }
+        `),
+        },
+        ROUTER_COMPATIBILITY_VERSION_ONE,
+      );
+      expect(costs.fieldWeights.get('Query.field')).toBeUndefined();
+    });
+
+    test('that a directive argument with a default value records directiveArgumentWeights even without explicit value', () => {
+      const { costs } = normalizeSubgraphSuccess(
+        {
+          name: 'subgraph-default-directive-arg',
+          url: '',
+          definitions: parse(`
+          directive @myDirective(arg1: Int = 1 @cost(weight: 5)) on FIELD_DEFINITION
+          type Query {
+            field: String! @myDirective
+          }
+        `),
+        },
+        ROUTER_COMPATIBILITY_VERSION_ONE,
+      );
+      expect(costs.fieldWeights.get('Query.field')).toEqual(
+        fieldWeight('Query', 'field', { directiveArgumentWeights: new Map([['myDirective.arg1', 5]]) }),
+      );
+    });
+
     test('that costs without directive argument weights has empty directiveArgumentWeights', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnField, ROUTER_COMPATIBILITY_VERSION_ONE);
       expect(costs.directiveArgumentWeights.size).toBe(0);
@@ -551,31 +608,21 @@ describe('@cost directive tests', () => {
   describe('costs internal structure tests', () => {
     test('that @cost on a field populates fieldWeights correctly', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnField, ROUTER_COMPATIBILITY_VERSION_ONE);
-      expect(costs.fieldWeights.get('Query.expensiveField')).toEqual({
-        typeName: 'Query',
-        fieldName: 'expensiveField',
-        argumentWeights: new Map(),
-        weight: 10,
-      });
+      expect(costs.fieldWeights.get('Query.expensiveField')).toEqual(
+        fieldWeight('Query', 'expensiveField', { weight: 10 }),
+      );
     });
 
     test('that @cost on a field argument populates fieldWeights.argumentWeights', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnArgument, ROUTER_COMPATIBILITY_VERSION_ONE);
-      expect(costs.fieldWeights.get('Query.search')).toEqual({
-        typeName: 'Query',
-        fieldName: 'search',
-        argumentWeights: new Map([['query', 5]]),
-      });
+      expect(costs.fieldWeights.get('Query.search')).toEqual(
+        fieldWeight('Query', 'search', { argumentWeights: new Map([['query', 5]]) }),
+      );
     });
 
     test('that @cost on an input field populates fieldWeights correctly', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnInputField, ROUTER_COMPATIBILITY_VERSION_ONE);
-      expect(costs.fieldWeights.get('SearchInput.query')).toEqual({
-        typeName: 'SearchInput',
-        fieldName: 'query',
-        argumentWeights: new Map(),
-        weight: 5,
-      });
+      expect(costs.fieldWeights.get('SearchInput.query')).toEqual(fieldWeight('SearchInput', 'query', { weight: 5 }));
     });
 
     test('that @cost on an object type populates typeWeights', () => {
@@ -610,12 +657,9 @@ describe('@cost directive tests', () => {
 
     test('that @cost on both a field and its argument populates a single FieldWeightConfiguration', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnFieldAndArgument, ROUTER_COMPATIBILITY_VERSION_ONE);
-      expect(costs.fieldWeights.get('Query.search')).toEqual({
-        typeName: 'Query',
-        fieldName: 'search',
-        weight: 10,
-        argumentWeights: new Map([['query', 3]]),
-      });
+      expect(costs.fieldWeights.get('Query.search')).toEqual(
+        fieldWeight('Query', 'search', { weight: 10, argumentWeights: new Map([['query', 3]]) }),
+      );
     });
 
     test('that a subgraph without cost directives has empty costs', () => {
@@ -630,12 +674,7 @@ describe('@cost directive tests', () => {
   describe('extension type tests', () => {
     test('that @cost on a field in an extend type populates fieldWeights', () => {
       const { costs } = normalizeSubgraphSuccess(subgraphWithCostOnExtensionField, ROUTER_COMPATIBILITY_VERSION_ONE);
-      expect(costs.fieldWeights.get('User.name')).toEqual({
-        typeName: 'User',
-        fieldName: 'name',
-        argumentWeights: new Map(),
-        weight: 5,
-      });
+      expect(costs.fieldWeights.get('User.name')).toEqual(fieldWeight('User', 'name', { weight: 5 }));
     });
 
     test('that @cost on an extend type populates typeWeights', () => {