Fix incomplete source generation of union types in initial render

workspaces/ballerina/ballerina-visualizer/src/views/BI/HelperPaneNew/Components/RecordConstructView/Types/UnionType/index.tsx

@@ -40,14 +40,70 @@ export default function UnionType(props: TypeProps) {
     const [paramSelected, setParamSelected] = useState(param.selected || requiredParam);
     const [selectedMemberType, setSelectedMemberType] = useState(getUnionParamName(initSelectedMember));
     const [parameter, setParameter] = useState<TypeField>(initSelectedMember);
+    const isInitialized = useRef(false);
 
-    // Initialize: If the union is selected, ensure the selected member and its required fields are also selected
+    // Initialize union member selection when param becomes selected
+    const initializeUnionMember = () => {
+        if (!param.members || param.members.length === 0) {
+            return false;
+        }
+
+        // Check if a member is already selected
+        const hasSelectedMember = param.members.some(member => member.selected === true);
+        if (hasSelectedMember) {
+            return true; // Already initialized
+        }
+
+        // Get the member to select (use initSelectedMember if available, otherwise first member)
+        const memberToSelect = initSelectedMember || param.members[0];
+
+        if (memberToSelect) {
+            const memberTypeName = getUnionParamName(memberToSelect);
+            if (memberTypeName) {
+                const selectedMember = param.members.find((field) => getUnionParamName(field) === memberTypeName);
+
+                if (selectedMember) {
+                    updateFormFieldMemberSelection(selectedMember);
+
+                    if (selectedMember.fields && selectedMember.fields.length > 0) {
+                        updateFieldsSelection(selectedMember.fields, true);
+                    }
+
+                    setSelectedMemberType(memberTypeName);
+                    setParameter(selectedMember);
+
+                    return true;
+                }
+            }
+        }
+
+        return false; // Failed to initialize
+    };
+
+    // Initialize: Always ensure a member is selected if union param is selected
     useEffect(() => {
-        if (paramSelected && initSelectedMember) {
-            handleMemberType(paramSelected ? selectedMemberType : "", false);
+        // Only run initialization once
+        if (isInitialized.current) {
+            return;
+        }
+        isInitialized.current = true;
+
+        // If union param is selected (or required), ensure a member is selected
+        if (paramSelected && param.members && param.members.length > 0) {
+            initializeUnionMember();
         }
     }, []);
 
+    useEffect(() => {
+        // If param becomes selected but we haven't initialized the member yet
+        if (param.selected && !isInitialized.current && param.members && param.members.length > 0) {
+            const initialized = initializeUnionMember();
+            if (initialized) {
+                isInitialized.current = true;
+            }
+        }
+    }, [param.selected]);
+
     if (!(param.members && param.members.length > 0)) {
         return <></>;
     }
@@ -69,8 +125,28 @@ export default function UnionType(props: TypeProps) {
     };
 
     const handleMemberType = (type: string, inCheckboxTrigger: boolean = true) => {
-        const selectedMember = param.members.find((field) => getUnionParamName(field) === type);
+        if (!type) {
+            return;
+        }
+
+        const selectedMember = param.members?.find((field) => getUnionParamName(field) === type);
+
+        if (!selectedMember) {
+            return;
+        }
+
+        // Ensure the union param itself is selected when selecting a member
+        if (!param.selected) {
+            param.selected = true;
+            setParamSelected(true);
+        }
+
         updateFormFieldMemberSelection(selectedMember);
+
+        if (selectedMember.fields && selectedMember.fields.length > 0) {
+            updateFieldsSelection(selectedMember.fields, true);
+        }
+
         setSelectedMemberType(type);
         setParameter(selectedMember);
 
@@ -82,11 +158,7 @@ export default function UnionType(props: TypeProps) {
             }
         }
 
-        // If the parent is selected and the selected member has fields, ensure required fields are selected
-        if (param.selected && selectedMember && selectedMember.fields && selectedMember.fields.length > 0) {
-            updateFieldsSelection(selectedMember.fields, true);
-        }
-
+        // Call onChange only after all selections are complete
         onChange();
     };
 
@@ -96,6 +168,12 @@ export default function UnionType(props: TypeProps) {
 
         // When checkbox is checked, ensure the currently selected member is also marked as selected
         if (newSelectedState) {
+            // Ensure member is initialized before proceeding
+            if (!isInitialized.current) {
+                initializeUnionMember();
+                isInitialized.current = true;
+            }
+
             const selectedMember = param.members.find((field) => getUnionParamName(field) === selectedMemberType);
             if (selectedMember) {
                 updateFormFieldMemberSelection(selectedMember);

workspaces/ballerina/ballerina-visualizer/src/views/BI/HelperPaneNew/Components/RecordConstructView/utils/index.tsx

@@ -72,9 +72,39 @@ export function updateFieldsSelection(fields: FormField[], selected: boolean): v
         // When deselecting: deselect all fields (both required and optional)
         if (!selected || isRequiredParam(field)) {
             field.selected = selected;
+
+            // If selecting a union type field, ensure a member is selected synchronously
+            if (selected && (field.typeName === "union" || field.typeName === "enum") && field.members && field.members.length > 0) {
+                // Check if a member is already selected
+                const hasSelectedMember = field.members.some(member => member.selected === true);
+
+                if (!hasSelectedMember) {
+                    // Select the first member (or use getSelectedUnionMember logic)
+                    const memberToSelect = getSelectedUnionMember(field) || field.members[0];
+                    if (memberToSelect) {
+                        const memberName = getUnionFormFieldName(memberToSelect);
+                        // Mark the selected member
+                        field.members.forEach(member => {
+                            member.selected = getUnionFormFieldName(member) === memberName;
+
+                            // If this is the selected member and it has nested fields, select required fields
+                            if (member.selected && member.fields && member.fields.length > 0) {
+                                updateFieldsSelection(member.fields, true);
+                            }
+                        });
+                    }
+                } else {
+                    // If a member is already selected, ensure its required fields are selected
+                    const selectedMember = field.members.find(member => member.selected === true);
+                    if (selectedMember && selectedMember.fields && selectedMember.fields.length > 0) {
+                        updateFieldsSelection(selectedMember.fields, true);
+                    }
+                }
+            }
         }
 
         // Recursively process nested fields
+        // Note: For union types, we handle members above, but union members can have fields too
         if (field.fields && field.fields.length > 0) {
             updateFieldsSelection(field.fields, selected);
         }