feat: support costs on arguments of directives

execution/engine/execution_engine_cost_test.go

@@ -14,9 +14,10 @@ func TestExecutionEngine_Cost(t *testing.T) {
 
 	t.Run("common on star wars scheme", func(t *testing.T) {
 		rootNodes := []plan.TypeField{
-			{TypeName: "Query", FieldNames: []string{"hero", "droid"}},
+			{TypeName: "Query", FieldNames: []string{"hero", "droid", "search"}},
 			{TypeName: "Human", FieldNames: []string{"name", "height", "friends"}},
 			{TypeName: "Droid", FieldNames: []string{"name", "primaryFunction", "friends"}},
+			{TypeName: "Starship", FieldNames: []string{"name", "length"}},
 		}
 		childNodes := []plan.TypeField{
 			{TypeName: "Character", FieldNames: []string{"name", "friends"}},
@@ -712,6 +713,127 @@ func TestExecutionEngine_Cost(t *testing.T) {
 			computeCosts(),
 		))
 
+		t.Run("field with @approx directive - cost weight on directive argument", runWithoutError(
+			ExecutionEngineTestCase{
+				schema: graphql.StarwarsSchema(t),
+				operation: func(t *testing.T) graphql.Request {
+					return graphql.Request{
+						Query: `{
+								search(name: "Luke") {
+									... on Human { name }
+								}
+							}`,
+					}
+				},
+				dataSources: []plan.DataSource{
+					mustGraphqlDataSourceConfiguration(t, "id",
+						mustFactory(t,
+							testNetHttpClient(t, roundTripperTestCase{
+								expectedHost:     "example.com",
+								expectedPath:     "/",
+								expectedBody:     "",
+								sendResponseBody: `{"data":{"search":{"__typename":"Human","name":"Luke"}}}`,
+								sendStatusCode:   200,
+							}),
+						),
+						&plan.DataSourceMetadata{
+							RootNodes:  rootNodes,
+							ChildNodes: childNodes,
+							CostConfig: &plan.DataSourceCostConfig{
+								Weights: map[plan.FieldCoordinate]*plan.FieldWeight{
+									{TypeName: "Query", FieldName: "search"}: {
+										HasWeight:       true,
+										Weight:          3,
+										ArgumentWeights: map[string]int{"name": 2},
+									},
+									{TypeName: "Human", FieldName: "name"}: {HasWeight: true, Weight: 5},
+								},
+								DirectiveArguments: map[plan.DirectiveArgCoords]int{
+									{DirectiveName: "approx", ArgName: "tolerance"}: -5,
+									// @deprecated is not on this field so it should not contribute.
+									{DirectiveName: "deprecated", ArgName: "reason"}: -100,
+								},
+							},
+						},
+						customConfig,
+					),
+				},
+				fields: []plan.FieldConfiguration{
+					{
+						TypeName: "Query", FieldName: "search",
+						Arguments: []plan.ArgumentConfiguration{
+							{
+								Name:         "name",
+								SourceType:   plan.FieldArgumentSource,
+								RenderConfig: plan.RenderArgumentAsGraphQLValue,
+							},
+						},
+					},
+				},
+				expectedResponse: `{"data":{"search":{"name":"Luke"}}}`,
+				// Query.search(3) + name arg(2) + Human.name(5) + @approx.tolerance(-5) = 5
+				expectedEstimatedCost: intPtr(5),
+				expectedActualCost:    intPtr(5),
+			},
+			computeCosts(),
+		))
+
+		t.Run("field with null directive arg does not affect cost", runWithoutError(
+			ExecutionEngineTestCase{
+				schema: graphql.StarwarsSchema(t),
+				operation: func(t *testing.T) graphql.Request {
+					return graphql.Request{
+						Query: `{
+								droid(id: "R2D2") {
+									name
+									primaryFunction
+								}
+							}`,
+					}
+				},
+				dataSources: []plan.DataSource{
+					mustGraphqlDataSourceConfiguration(t, "id",
+						mustFactory(t,
+							testNetHttpClient(t, roundTripperTestCase{
+								expectedHost: "example.com", expectedPath: "/", expectedBody: "",
+								sendResponseBody: `{"data":{"droid":{"name":"R2D2","primaryFunction":"no"}}}`,
+								sendStatusCode:   200,
+							}),
+						),
+						&plan.DataSourceMetadata{
+							RootNodes:  rootNodes,
+							ChildNodes: childNodes,
+							CostConfig: &plan.DataSourceCostConfig{
+								Weights: map[plan.FieldCoordinate]*plan.FieldWeight{
+									{TypeName: "Droid", FieldName: "name"}: {HasWeight: true, Weight: 17},
+								},
+								DirectiveArguments: map[plan.DirectiveArgCoords]int{
+									{DirectiveName: "approx", ArgName: "tolerance"}: -5,
+								},
+							}},
+						customConfig,
+					),
+				},
+				fields: []plan.FieldConfiguration{
+					{
+						TypeName: "Query", FieldName: "droid",
+						Arguments: []plan.ArgumentConfiguration{
+							{
+								Name:         "id",
+								SourceType:   plan.FieldArgumentSource,
+								RenderConfig: plan.RenderArgumentAsGraphQLValue,
+							},
+						},
+					},
+				},
+				expectedResponse: `{"data":{"droid":{"name":"R2D2","primaryFunction":"no"}}}`,
+				// Query.droid (1) + droid.name (17); @approx.tolerance is null
+				expectedEstimatedCost: intPtr(18),
+				expectedActualCost:    intPtr(18),
+			},
+			computeCosts(),
+		))
+
 	})
 
 	t.Run("union types", func(t *testing.T) {

v2/pkg/engine/plan/cost.go

@@ -113,6 +113,11 @@ func (ls *FieldListSize) multiplier(arguments map[string]ArgumentInfo, vars *ast
 	return multiplier
 }
 
+type DirectiveArgCoords struct {
+	DirectiveName string
+	ArgName       string
+}
+
 // DataSourceCostConfig holds all cost configurations for a data source.
 // This data is passed from the composition.
 type DataSourceCostConfig struct {
@@ -130,12 +135,11 @@ type DataSourceCostConfig struct {
 	// Location: ENUM, OBJECT, SCALAR
 	Types map[string]int
 
-	// Arguments on directives is a special case. They use a special kind of coordinate:
-	// directive name + argument name. That should be the key mapped to the weight.
-	//
-	// Directives can be used on [input] object fields and arguments of fields. This creates
+	// DirectiveArguments maps Arguments on directives to the weight.
+	// Directives can be used on input object fields and arguments of fields. This creates
 	// mutual recursion between them; it complicates cost calculation.
-	// We avoid them intentionally in the first iteration.
+	// We avoid them intentionally for now.
+	DirectiveArguments map[DirectiveArgCoords]int
 }
 
 // NewDataSourceCostConfig creates a new cost config with defaults
@@ -198,6 +202,9 @@ type CostTreeNode struct {
 	// arguments contain the values of arguments passed to the field.
 	arguments map[string]ArgumentInfo
 
+	// directiveArguments contain directive arguments of non-null value for this field.
+	directiveArguments []DirectiveArgCoords
+
 	jsonPath string // JSON path using aliases too
 
 	returnsListType         bool
@@ -229,6 +236,16 @@ type ArgumentInfo struct {
 	varName string
 }
 
+// DirectiveInfo stores information about a directive applied to a field definition,
+// including its name and the arguments it was invoked with.
+type DirectiveInfo struct {
+	// name of the directive (e.g. "approx" for @approx)
+	name string
+
+	// arguments indicate the presence of a non-null value for an argument name
+	arguments map[string]struct{}
+}
+
 // inputObjectField describes the type of input object field.
 type inputObjectField struct {
 	unwrappedTypeName string
@@ -479,6 +496,13 @@ func (node *CostTreeNode) costsAndMultiplier(configs map[DSHash]*DataSourceCostC
 			}
 		}
 
+		for _, dirArgCoords := range node.directiveArguments {
+			weight, ok := dsCostConfig.DirectiveArguments[dirArgCoords]
+			if ok {
+				directiveCost += weight
+			}
+		}
+
 		if !node.returnsListType || !isEstimation {
 			continue
 		}

v2/pkg/engine/plan/cost_visitor.go

@@ -66,6 +66,7 @@ func (v *CostVisitor) EnterField(fieldRef int) {
 	unwrappedTypeName := v.Definition.ResolveTypeNameString(fieldDefinitionTypeRef)
 
 	arguments := v.extractFieldArguments(fieldRef)
+	directiveArgs := v.extractFieldDirectives(fieldDefinitionRef)
 
 	// Check and push through if the unwrapped type of this field is interface or union.
 	unwrappedTypeNode, exists := v.Definition.NodeByNameStr(unwrappedTypeName)
@@ -101,7 +102,7 @@ func (v *CostVisitor) EnterField(fieldRef int) {
 	}
 
 	isEnclosingTypeAbstract := v.Walker.EnclosingTypeDefinition.Kind.IsAbstractType()
-	// Create a skeleton node. dataSourceHashes will be filled in leaveFieldCost
+	// Partially filled node. dataSourceHashes will be filled in leaveFieldCost
 	node := CostTreeNode{
 		fieldRef:                fieldRef,
 		fieldCoords:             FieldCoordinate{typeName, fieldName},
@@ -112,10 +113,11 @@ func (v *CostVisitor) EnterField(fieldRef int) {
 		returnsAbstractType:     isAbstractType,
 		isEnclosingTypeAbstract: isEnclosingTypeAbstract,
 		arguments:               arguments,
+		directiveArguments:      directiveArgs,
 		jsonPath:                jsonPath,
 	}
 
-	// Attach to parent
+	// Attach to the parent
 	if len(v.stack) > 0 {
 		parent := v.stack[len(v.stack)-1]
 		parent.children = append(parent.children, &node)
@@ -265,6 +267,57 @@ func (v *CostVisitor) buildInputObjectFieldTypes(typeName string, node ast.Node,
 	}
 }
 
+// extractFieldDirectives extracts directives applied to a field definition in the schema.
+func (v *CostVisitor) extractFieldDirectives(fieldDefinitionRef int) []DirectiveArgCoords {
+	directiveRefs := v.Definition.FieldDefinitionDirectives(fieldDefinitionRef)
+	if len(directiveRefs) == 0 {
+		return nil
+	}
+
+	directives := make([]DirectiveArgCoords, 0, len(directiveRefs))
+	for _, directiveRef := range directiveRefs {
+		directiveName := v.Definition.DirectiveNameString(directiveRef)
+		args := make(map[string]struct{})
+
+		// Populate from the directive definition to make defaulted arguments visible
+		// during cost calculation.
+		directiveDefRef, ok := v.Definition.DirectiveDefinitionByName(directiveName)
+		if !ok {
+			continue
+		}
+		directiveDef := v.Definition.DirectiveDefinitions[directiveDefRef]
+		if directiveDef.HasArgumentsDefinitions {
+			for _, inputValRef := range directiveDef.ArgumentsDefinition.Refs {
+				argName := v.Definition.InputValueDefinitionNameString(inputValRef)
+				defaultValue := v.Definition.InputValueDefinitionDefaultValue(inputValRef)
+				if defaultValue.Kind != ast.ValueKindUnknown && defaultValue.Kind != ast.ValueKindNull {
+					args[argName] = struct{}{}
+				}
+			}
+		}
+
+		// Override with explicitly provided arguments at the usage site.
+		// Null values exclude the argument from cost calculation (even if a default existed).
+		for _, argRef := range v.Definition.DirectiveArgumentSet(directiveRef) {
+			argName := v.Definition.ArgumentNameString(argRef)
+			if v.Definition.ArgumentValue(argRef).Kind == ast.ValueKindNull {
+				delete(args, argName)
+			} else {
+				args[argName] = struct{}{}
+			}
+		}
+
+		for argName := range args {
+			directives = append(directives, DirectiveArgCoords{
+				DirectiveName: directiveName,
+				ArgName:       argName,
+			})
+		}
+	}
+
+	return directives
+}
+
 func (v *CostVisitor) finalCostTree() *CostTreeNode {
 	return v.tree
 }

v2/pkg/starwars/testdata/star_wars.graphql

@@ -8,10 +8,13 @@ schema {
 
 directive @testDeprecated(okArg: String deprecatedArg: String @deprecated(reason: "no such arg")) on FIELD_DEFINITION
 
+# the tolerance argument has "@cost(weight: -5)" defined in tests
+directive @approx(tolerance: Int = 1) on FIELD_DEFINITION
+
 type Query {
     hero: Character @deprecated
-    droid(id: ID!): Droid
-    search(name: String!): SearchResult
+    droid(id: ID!): Droid @approx(tolerance: null)
+    search(name: String!): SearchResult @approx
     searchResults: [SearchResult]
 }