diff --git a/app/services/namespaces/projects/flows/flow_service_helper.rb b/app/services/namespaces/projects/flows/flow_service_helper.rb
index 8efc0d03..ef7cd91b 100644
--- a/app/services/namespaces/projects/flows/flow_service_helper.rb
+++ b/app/services/namespaces/projects/flows/flow_service_helper.rb
@@ -11,7 +11,7 @@ def get_data_type_identifier(runtime, identifier, t)
           end
 
           if identifier.generic_type.present?
-            data_type = namespace_project.primary_runtime.data_types.find_by(
+            data_type = runtime.data_types.find_by(
               id: identifier.generic_type.data_type_id.model_id
             )
 
@@ -35,7 +35,7 @@ def get_data_type_identifier(runtime, identifier, t)
             return DataTypeIdentifier.find_or_create_by(runtime: runtime, generic_type: generic_type)
           end
 
-          data_type = namespace_project.primary_runtime.data_types.find_by(id: identifier.data_type_id.model_id)
+          data_type = runtime.data_types.find_by(id: identifier.data_type_id.model_id)
 
           if data_type.nil?
             t.rollback_and_return! ServiceResponse.error(
diff --git a/spec/requests/graphql/mutation/namespace/projects/flows/update_mutation_spec.rb b/spec/requests/graphql/mutation/namespace/projects/flows/update_mutation_spec.rb
index 9997c7ee..c1c014cd 100644
--- a/spec/requests/graphql/mutation/namespace/projects/flows/update_mutation_spec.rb
+++ b/spec/requests/graphql/mutation/namespace/projects/flows/update_mutation_spec.rb
@@ -72,7 +72,8 @@
       runtime_function_definition: rfd,
       data_type: create(
         :data_type_identifier,
-        data_type: create(:data_type)
+        runtime: runtime,
+        data_type: create(:data_type, runtime: runtime)
       )
     )
 
@@ -142,7 +143,7 @@
                   ],
                   nodeFunctionId: 'gid://sagittarius/NodeFunction/2000',
                   dataTypeIdentifier: {
-                    genericKey: 'K',
+                    dataTypeId: function_definition.parameter_definitions.first.data_type.data_type.to_global_id,
                   },
                 },
               }