STET

agent_patch

runs/2026-02-21__03-36-16__gpt-5-1-codex-mini/flux-pr-1262/flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/agent-logs/agent.patch

19565 bytes

Back to adjudication
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go b/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
index 6ee1940..93445a2 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/configuration.go
@@ -17,7 +17,7 @@ type GRPCMapping struct {
 	// SubscriptionRPCs maps GraphQL subscription fields to the corresponding gRPC RPC configurations
 	SubscriptionRPCs RPCConfigMap
 	// EntityRPCs defines how GraphQL types are resolved as entities using specific RPCs
-	EntityRPCs map[string]EntityRPCConfig
+	EntityRPCs map[string][]EntityRPCConfig
 	// Fields defines the field mappings between GraphQL types and gRPC messages
 	Fields map[string]FieldMap
 	// EnumValues defines the enum values for each enum type
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
index 5756d9d..ea11505 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_federation_test.go
@@ -21,13 +21,15 @@ func TestEntityLookup(t *testing.T) {
 			query: `query EntityLookup($representations: [_Any!]!) { _entities(representations: $representations) { ... on Product { __typename id name price } } }`,
 			mapping: &GRPCMapping{
 				Service: "Products",
-				EntityRPCs: map[string]EntityRPCConfig{
+				EntityRPCs: map[string][]EntityRPCConfig{
 					"Product": {
-						Key: "id",
-						RPCConfig: RPCConfig{
-							RPC:      "LookupProductById",
-							Request:  "LookupProductByIdRequest",
-							Response: "LookupProductByIdResponse",
+						{
+							Key: "id",
+							RPCConfig: RPCConfig{
+								RPC:      "LookupProductById",
+								Request:  "LookupProductByIdRequest",
+								Response: "LookupProductByIdResponse",
+							},
 						},
 					},
 				},
@@ -49,6 +51,11 @@ func TestEntityLookup(t *testing.T) {
 									Message: &RPCMessage{
 										Name: "LookupProductByIdKey",
 										Fields: []RPCField{
+											{
+												Name:     "__typename",
+												TypeName: string(DataTypeString),
+												JSONPath: "__typename",
+											},
 											{
 												Name:     "id",
 												TypeName: string(DataTypeString),
@@ -69,28 +76,174 @@ func TestEntityLookup(t *testing.T) {
 									Repeated: true,
 									JSONPath: "_entities",
 									Message: &RPCMessage{
-										Name: "Product",
+										Name:   "_Entity",
+										Fields: nil,
+										FieldSelectionSet: RPCFieldSelectionSet{
+											"Product": {
+												{
+													Name:        "__typename",
+													TypeName:    string(DataTypeString),
+													JSONPath:    "__typename",
+													StaticValue: "Product",
+												},
+												{
+													Name:     "id",
+													TypeName: string(DataTypeString),
+													JSONPath: "id",
+												},
+												{
+													Name:     "name",
+													TypeName: string(DataTypeString),
+													JSONPath: "name",
+												},
+												{
+													Name:     "price",
+													TypeName: string(DataTypeDouble),
+													JSONPath: "price",
+												},
+											},
+										},
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "Should create an execution plan for an entity lookup with multiple key types",
+			query: `
+query EntityLookup($representations: [_Any!]!) {
+	_entities(representations: $representations) {
+		... on Product {
+			__typename
+			id
+			name
+			price
+		}
+		... on Storage {
+			__typename
+			id
+			name
+			location
+		}
+	}
+}
+`,
+			mapping: &GRPCMapping{
+				Service: "Products",
+				EntityRPCs: map[string][]EntityRPCConfig{
+					"Product": {
+						{
+							Key: "id",
+							RPCConfig: RPCConfig{
+								RPC:      "LookupProductById",
+								Request:  "LookupProductByIdRequest",
+								Response: "LookupProductByIdResponse",
+							},
+						},
+					},
+					"Storage": {
+						{
+							Key: "id",
+							RPCConfig: RPCConfig{
+								RPC:      "LookupStorageById",
+								Request:  "LookupStorageByIdRequest",
+								Response: "LookupStorageByIdResponse",
+							},
+						},
+					},
+				},
+			},
+			expectedPlan: &RPCExecutionPlan{
+				Calls: []RPCCall{
+					{
+						ServiceName: "Products",
+						MethodName:  "LookupProductById",
+						Request: RPCMessage{
+							Name: "LookupProductByIdRequest",
+							Fields: []RPCField{
+								{
+									Name:     "keys",
+									TypeName: string(DataTypeMessage),
+									Repeated: true,
+									JSONPath: "representations",
+									Message: &RPCMessage{
+										Name: "LookupProductByIdKey",
 										Fields: []RPCField{
 											{
-												Name:        "__typename",
-												TypeName:    string(DataTypeString),
-												JSONPath:    "__typename",
-												StaticValue: "Product",
+												Name:     "__typename",
+												TypeName: string(DataTypeString),
+												JSONPath: "__typename",
 											},
 											{
 												Name:     "id",
 												TypeName: string(DataTypeString),
 												JSONPath: "id",
 											},
-											{
-												Name:     "name",
-												TypeName: string(DataTypeString),
-												JSONPath: "name",
+										},
+									},
+								},
+							},
+						},
+						Response: RPCMessage{
+							Name: "LookupProductByIdResponse",
+							Fields: []RPCField{
+								{
+									Name:     "result",
+									TypeName: string(DataTypeMessage),
+									Repeated: true,
+									JSONPath: "_entities",
+									Message: &RPCMessage{
+										Name:   "_Entity",
+										Fields: nil,
+										FieldSelectionSet: RPCFieldSelectionSet{
+											"Product": {
+												{
+													Name:        "__typename",
+													TypeName:    string(DataTypeString),
+													JSONPath:    "__typename",
+													StaticValue: "Product",
+												},
+												{
+													Name:     "id",
+													TypeName: string(DataTypeString),
+													JSONPath: "id",
+												},
+												{
+													Name:     "name",
+													TypeName: string(DataTypeString),
+													JSONPath: "name",
+												},
+												{
+													Name:     "price",
+													TypeName: string(DataTypeDouble),
+													JSONPath: "price",
+												},
 											},
-											{
-												Name:     "price",
-												TypeName: string(DataTypeDouble),
-												JSONPath: "price",
+											"Storage": {
+												{
+													Name:        "__typename",
+													TypeName:    string(DataTypeString),
+													JSONPath:    "__typename",
+													StaticValue: "Storage",
+												},
+												{
+													Name:     "id",
+													TypeName: string(DataTypeString),
+													JSONPath: "id",
+												},
+												{
+													Name:     "name",
+													TypeName: string(DataTypeString),
+													JSONPath: "name",
+												},
+												{
+													Name:     "location",
+													TypeName: string(DataTypeString),
+													JSONPath: "location",
+												},
 											},
 										},
 									},
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
index edf2dd8..858442f 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/execution_plan_visitor.go
@@ -17,11 +17,48 @@ type keyField struct {
 }
 
 type entityInfo struct {
-	name                    string
-	keyFields               []keyField
-	keyTypeName             string
-	entityRootFieldRef      int
-	entityInlineFragmentRef int
+	name               string
+	keyFields          []keyField
+	keyFieldsSet       map[string]struct{}
+	keyTypeName        string
+	entityRootFieldRef int
+	scaffolded         bool
+	keyMessage         *RPCMessage
+}
+
+func (e *entityInfo) addKeyField(field keyField) bool {
+	if e.keyFieldsSet == nil {
+		e.keyFieldsSet = make(map[string]struct{})
+	}
+	if _, exists := e.keyFieldsSet[field.fieldName]; exists {
+		return false
+	}
+	e.keyFieldsSet[field.fieldName] = struct{}{}
+	e.keyFields = append(e.keyFields, field)
+	e.updateKeyTypeName()
+	return true
+}
+
+func (e *entityInfo) updateKeyTypeName() {
+	if e.name == "" || len(e.keyFields) == 0 {
+		return
+	}
+	fieldNames := make([]string, len(e.keyFields))
+	for i, key := range e.keyFields {
+		fieldNames[i] = key.fieldName
+	}
+	e.keyTypeName = e.name + "By" + strings.Join(titleSlice(fieldNames), "And")
+}
+
+func (e *entityInfo) addKeyFieldToMessage(field keyField) {
+	if e.keyMessage == nil {
+		return
+	}
+	e.keyMessage.Fields = append(e.keyMessage.Fields, RPCField{
+		Name:     field.fieldName,
+		TypeName: field.fieldType,
+		JSONPath: field.fieldName,
+	})
 }
 
 type planningInfo struct {
@@ -233,20 +270,18 @@ func (r *rpcPlanVisitor) LeaveSelectionSet(ref int) {
 // EnterInlineFragment implements astvisitor.InlineFragmentVisitor.
 func (r *rpcPlanVisitor) EnterInlineFragment(ref int) {
 	entityInfo := &r.planInfo.entityInfo
-	if entityInfo.entityRootFieldRef != -1 && entityInfo.entityInlineFragmentRef == -1 {
-		entityInfo.entityInlineFragmentRef = ref
-		r.resolveEntityInformation(ref)
-		r.scaffoldEntityLookup()
-
+	if entityInfo.entityRootFieldRef == -1 {
 		return
 	}
+	if !entityInfo.scaffolded {
+		r.scaffoldEntityLookup()
+	}
+	r.resolveEntityInformation(ref)
 }
 
 // LeaveInlineFragment implements astvisitor.InlineFragmentVisitor.
 func (r *rpcPlanVisitor) LeaveInlineFragment(ref int) {
-	if ref == r.planInfo.entityInfo.entityInlineFragmentRef {
-		r.planInfo.entityInfo.entityInlineFragmentRef = -1
-	}
+	// no-op
 }
 
 func (r *rpcPlanVisitor) IsRootField() bool {
@@ -302,11 +337,9 @@ func (r *rpcPlanVisitor) EnterField(ref int) {
 		// Providing variables (_Any) and the response type is a Union that needs to be
 		// determined from the first inline fragment.
 		r.planInfo.entityInfo = entityInfo{
-			entityRootFieldRef:      ref,
-			entityInlineFragmentRef: -1,
+			entityRootFieldRef: ref,
 		}
 		r.planInfo.isEntityLookup = true
-		r.planInfo.entityInfo.entityRootFieldRef = ref
 		return
 	}
 
@@ -326,7 +359,7 @@ func (r *rpcPlanVisitor) EnterField(ref int) {
 
 	field := r.buildField(fd, fieldName, fieldAlias)
 
-	if ref, ok := r.IsInlineFragmentField(); ok && !r.planInfo.isEntityLookup {
+	if ref, ok := r.IsInlineFragmentField(); ok {
 		if r.planInfo.currentResponseMessage.FieldSelectionSet == nil {
 			r.planInfo.currentResponseMessage.FieldSelectionSet = make(RPCFieldSelectionSet)
 		}
@@ -585,79 +618,101 @@ func (r *rpcPlanVisitor) createListMetadata(typeRef int) *ListMetadata {
 }
 
 func (r *rpcPlanVisitor) resolveEntityInformation(inlineFragmentRef int) {
-	// TODO support multiple entities in a single query
-	if !r.planInfo.isEntityLookup || r.planInfo.entityInfo.name != "" {
+	if !r.planInfo.isEntityLookup {
 		return
 	}
 
 	fragmentName := r.operation.InlineFragmentTypeConditionNameString(inlineFragmentRef)
+	if fragmentName == "" {
+		return
+	}
+
 	node, found := r.definition.NodeByNameStr(fragmentName)
 	if !found {
 		return
 	}
 
-	// Only process object type definitions
-	// TODO: handle interfaces
 	if node.Kind != ast.NodeKindObjectTypeDefinition {
 		return
 	}
 
-	// An entity must at least have a key directive
 	def := r.definition.ObjectTypeDefinitions[node.Ref]
 	if !def.HasDirectives {
 		return
 	}
 
-	// TODO: We currently only support one key directive per entity
-	// We need to get the used key from the graphql datasource.
+	entityInfo := &r.planInfo.entityInfo
+	configureCall := entityInfo.name == ""
+	if configureCall {
+		entityInfo.name = fragmentName
+	}
+	var configs []EntityRPCConfig
+	var hasConfigs bool
+	var selectedConfig *EntityRPCConfig
+	if configureCall {
+		configs, hasConfigs = r.mapping.EntityRPCs[fragmentName]
+	}
+
 	for _, directiveRef := range def.Directives.Refs {
 		if r.definition.DirectiveNameString(directiveRef) != federationKeyDirectiveName {
 			continue
 		}
 
-		r.planInfo.entityInfo.name = fragmentName
-
 		directive := r.definition.Directives[directiveRef]
 		for _, argRef := range directive.Arguments.Refs {
 			if r.definition.ArgumentNameString(argRef) != "fields" {
 				continue
 			}
 			argument := r.definition.Arguments[argRef]
-			keyFieldName := r.definition.ValueContentString(argument.Value)
-
-			fieldDef, ok := r.definition.NodeFieldDefinitionByName(node, ast.ByteSlice(keyFieldName))
-			if !ok {
-				r.walker.Report.AddExternalError(operationreport.ExternalError{
-					Message: fmt.Sprintf("Field %s not found in definition", keyFieldName),
-				})
-				return
+			fieldsStr := r.definition.ValueContentString(argument.Value)
+			fieldNames := strings.Fields(fieldsStr)
+			if len(fieldNames) == 0 {
+				continue
 			}
 
-			fdt := r.definition.FieldDefinitionType(fieldDef)
-			ft := r.definition.Types[fdt]
+			if configureCall && hasConfigs && selectedConfig == nil {
+				for j := range configs {
+					if configs[j].Key == fieldsStr {
+						selectedConfig = &configs[j]
+						break
+					}
+				}
+			}
 
-			r.planInfo.entityInfo.keyFields =
-				append(r.planInfo.entityInfo.keyFields, keyField{
+			for _, keyFieldName := range fieldNames {
+				fieldDef, ok := r.definition.NodeFieldDefinitionByName(node, ast.ByteSlice(keyFieldName))
+				if !ok {
+					r.walker.Report.AddExternalError(operationreport.ExternalError{
+						Message: fmt.Sprintf("Field %s not found in definition", keyFieldName),
+					})
+					return
+				}
+
+				fdt := r.definition.FieldDefinitionType(fieldDef)
+				ft := r.definition.Types[fdt]
+				key := keyField{
 					fieldName: keyFieldName,
 					fieldType: r.toDataType(&ft).String(),
-				})
+				}
+				if entityInfo.addKeyField(key) {
+					entityInfo.addKeyFieldToMessage(key)
+				}
+			}
 		}
-
-		break
 	}
 
-	keyFields := make([]string, 0, len(r.planInfo.entityInfo.keyFields))
-	for _, key := range r.planInfo.entityInfo.keyFields {
-		keyFields = append(keyFields, key.fieldName)
+	if configureCall && selectedConfig == nil && hasConfigs && len(configs) > 0 {
+		selectedConfig = &configs[0]
 	}
 
-	if ei, exists := r.mapping.EntityRPCs[r.planInfo.entityInfo.name]; exists {
-		r.currentCall.Request.Name = ei.RPCConfig.Request
-		r.currentCall.Response.Name = ei.RPCConfig.Response
-		r.planInfo.methodName = ei.RPCConfig.RPC
+	if selectedConfig != nil {
+		r.currentCall.Request.Name = selectedConfig.RPCConfig.Request
+		r.currentCall.Response.Name = selectedConfig.RPCConfig.Response
+		r.planInfo.methodName = selectedConfig.RPCConfig.RPC
+	}
+	if entityInfo.keyMessage != nil {
+		entityInfo.keyMessage.Name = r.rpcMethodName() + "Key"
 	}
-
-	r.planInfo.entityInfo.keyTypeName = r.planInfo.entityInfo.name + "By" + strings.Join(titleSlice(keyFields), "And")
 }
 
 // scaffoldEntityLookup creates the entity lookup call structure
@@ -669,16 +724,20 @@ func (r *rpcPlanVisitor) scaffoldEntityLookup() {
 	}
 
 	entityInfo := &r.planInfo.entityInfo
+	if entityInfo.scaffolded {
+		return
+	}
+
 	keyFieldMessage := &RPCMessage{
 		Name: r.rpcMethodName() + "Key",
 	}
-	for _, key := range entityInfo.keyFields {
-		keyFieldMessage.Fields = append(keyFieldMessage.Fields, RPCField{
-			Name:     key.fieldName,
-			TypeName: key.fieldType,
-			JSONPath: key.fieldName,
-		})
-	}
+	keyFieldMessage.Fields = append(keyFieldMessage.Fields, RPCField{
+		Name:     "__typename",
+		TypeName: string(DataTypeString),
+		JSONPath: "__typename",
+	})
+	entityInfo.keyMessage = keyFieldMessage
+	entityInfo.scaffolded = true
 
 	r.planInfo.currentRequestMessage.Fields = []RPCField{
 		{
@@ -698,6 +757,9 @@ func (r *rpcPlanVisitor) scaffoldEntityLookup() {
 			TypeName: DataTypeMessage.String(),
 			JSONPath: "_entities",
 			Repeated: true,
+			Message: &RPCMessage{
+				Name: "_Entity",
+			},
 		},
 	}
 }
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go b/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
index 146f371..1ace3cd 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
+++ b/app/v2/pkg/engine/datasource/grpc_datasource/mapping_test_helper.go
@@ -218,21 +218,25 @@ func testMapping() *GRPCMapping {
 			},
 		},
 		SubscriptionRPCs: RPCConfigMap{},
-		EntityRPCs: map[string]EntityRPCConfig{
+		EntityRPCs: map[string][]EntityRPCConfig{
 			"Product": {
-				Key: "id",
-				RPCConfig: RPCConfig{
-					RPC:      "LookupProductById",
-					Request:  "LookupProductByIdRequest",
-					Response: "LookupProductByIdResponse",
+				{
+					Key: "id",
+					RPCConfig: RPCConfig{
+						RPC:      "LookupProductById",
+						Request:  "LookupProductByIdRequest",
+						Response: "LookupProductByIdResponse",
+					},
 				},
 			},
 			"Storage": {
-				Key: "id",
-				RPCConfig: RPCConfig{
-					RPC:      "LookupStorageById",
-					Request:  "LookupStorageByIdRequest",
-					Response: "LookupStorageByIdResponse",
+				{
+					Key: "id",
+					RPCConfig: RPCConfig{
+						RPC:      "LookupStorageById",
+						Request:  "LookupStorageByIdRequest",
+						Response: "LookupStorageByIdResponse",
+					},
 				},
 			},
 		},
diff --git a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/grpctest/mapping/mapping.go b/app/v2/pkg/grpctest/mapping/mapping.go
index 842a029..1476d08 100644
--- a/tmp/agent-patch-flux-pr-1262.1-of-1.2026-03-05__17-48-26__gpt-5-1-codex-mini/app/v2/pkg/grpctest/mapping/mapping.go
+++ b/app/v2/pkg/grpctest/mapping/mapping.go
@@ -225,21 +225,25 @@ func DefaultGRPCMapping() *grpcdatasource.GRPCMapping {
 			},
 		},
 		SubscriptionRPCs: grpcdatasource.RPCConfigMap{},
-		EntityRPCs: map[string]grpcdatasource.EntityRPCConfig{
+		EntityRPCs: map[string][]grpcdatasource.EntityRPCConfig{
 			"Product": {
-				Key: "id",
-				RPCConfig: grpcdatasource.RPCConfig{
-					RPC:      "LookupProductById",
-					Request:  "LookupProductByIdRequest",
-					Response: "LookupProductByIdResponse",
+				{
+					Key: "id",
+					RPCConfig: grpcdatasource.RPCConfig{
+						RPC:      "LookupProductById",
+						Request:  "LookupProductByIdRequest",
+						Response: "LookupProductByIdResponse",
+					},
 				},
 			},
 			"Storage": {
-				Key: "id",
-				RPCConfig: grpcdatasource.RPCConfig{
-					RPC:      "LookupStorageById",
-					Request:  "LookupStorageByIdRequest",
-					Response: "LookupStorageByIdResponse",
+				{
+					Key: "id",
+					RPCConfig: grpcdatasource.RPCConfig{
+						RPC:      "LookupStorageById",
+						Request:  "LookupStorageByIdRequest",
+						Response: "LookupStorageByIdResponse",
+					},
 				},
 			},
 		},