Handle union types in Format2 Schema

server/gx-workflow-ls-format2/src/schema/definitions.ts

@@ -297,11 +297,23 @@ export class FieldSchemaNode implements SchemaNode, IdMapper {
     if (this.canBeArray) {
       return this.getArrayItemTypeName() || "undefined";
     }
-    const mainType = this._allowedTypes.find((t) => t.typeName !== "null");
-    //TODO: this requires more logic... we cannot assume the first non-null type to be the main
+    const nonNullTypes = this._allowedTypes.filter((t) => t.typeName !== "null");
+    if (nonNullTypes.length > 1) {
+      // Union type
+      return nonNullTypes.map((t) => t.typeName).join("|");
+    }
+    const mainType = nonNullTypes[0];
     return isBasicFieldType(mainType) ? mainType.typeName : "unknown";
   }
 
+  /**
+   * Returns the type references for union types or
+   * a single type reference for non-union types.
+   */
+  public get typeRefs(): string[] {
+    return this.typeRef.split("|");
+  }
+
   public get mapSubject(): string | undefined {
     return this._schemaField.jsonldPredicate?.mapSubject;
   }
@@ -316,14 +328,19 @@ export class FieldSchemaNode implements SchemaNode, IdMapper {
       return arrayType?.itemType.typeName;
     }
     if (arrayType?.itemType instanceof Array) {
-      return arrayType.itemType.map((i) => i.typeName).at(0); // TODO REMOVE AT
+      // Union type
+      return arrayType.itemType.map((t) => t.typeName).join("|");
     }
     console.debug("getArrayItemTypeName -> Type name not found");
     return undefined;
   }
 
   public get isPrimitiveType(): boolean {
-    return FieldSchemaNode.definitions.primitiveTypes.has(this.typeRef);
+    return FieldSchemaNode.definitions.isPrimitiveType(this.typeRef);
+  }
+
+  public get isUnionType(): boolean {
+    return this.typeRef.includes("|");
   }
 
   public get isObjectType(): boolean {
@@ -412,4 +429,5 @@ export interface SchemaDefinitions {
   enums: Map<string, EnumSchemaNode>;
   specializations: Map<string, string>;
   primitiveTypes: Set<string>;
+  isPrimitiveType(typeName: string): boolean;
 }

server/gx-workflow-ls-format2/src/schema/schemaLoader.ts

@@ -74,6 +74,9 @@ export class GalaxyWorkflowFormat2SchemaLoader implements GalaxyWorkflowSchemaLo
       enums: new Map<string, EnumSchemaNode>(),
       specializations: new Map<string, string>(),
       primitiveTypes: new Set<string>(),
+      isPrimitiveType: (type: string) => {
+        return definitions.primitiveTypes.has(type);
+      },
     };
 
     this.expandEntries(schemaEntries.values());

server/gx-workflow-ls-format2/src/schema/schemaNodeResolver.ts

@@ -32,8 +32,21 @@ export class SchemaNodeResolverImpl implements SchemaNodeResolver {
       if (currentSchemaNode instanceof RecordSchemaNode) {
         currentSchemaNode = currentSchemaNode.fields.find((f) => f.name === currentSegment);
       } else if (currentSchemaNode instanceof FieldSchemaNode) {
-        const typeNode = this.getSchemaNodeByTypeRef(currentSchemaNode.typeRef);
-        currentSchemaNode = typeNode;
+        if (currentSchemaNode.isUnionType) {
+          for (const typeRef of currentSchemaNode.typeRefs) {
+            const resolvedNode = this.getSchemaNodeByTypeRef(typeRef);
+            if (resolvedNode instanceof RecordSchemaNode) {
+              const matchedField = resolvedNode.fields.find((f) => f.name === currentSegment);
+              if (matchedField) {
+                currentSchemaNode = matchedField;
+                break;
+              }
+            }
+          }
+        } else {
+          const typeNode = this.getSchemaNodeByTypeRef(currentSchemaNode.typeRef);
+          currentSchemaNode = typeNode;
+        }
       }
       currentSegment = toWalk.shift();
     }

server/gx-workflow-ls-format2/src/services/completionService.ts

@@ -46,7 +46,25 @@ export class GxFormat2CompletionService {
     const overwriteRange = textBuffer.getCurrentWordRange(offset);
     const position = textBuffer.getPosition(offset);
     const isPositionAfterColon = textBuffer.isPositionAfterToken(position, ":");
-    if (schemaNode instanceof RecordSchemaNode) {
+    if (schemaNode instanceof EnumSchemaNode) {
+      schemaNode.symbols
+        .filter((v) => v.startsWith(currentWord))
+        .forEach((value) => {
+          if (exclude.has(value)) return;
+          const item: CompletionItem = {
+            label: value,
+            sortText: `_${value}`,
+            kind: CompletionItemKind.EnumMember,
+            documentation: schemaNode.documentation,
+            insertText: value,
+            textEdit: {
+              range: overwriteRange,
+              newText: value,
+            },
+          };
+          result.push(item);
+        });
+    } else if (schemaNode instanceof RecordSchemaNode) {
       if (isPositionAfterColon) {
         return result; // Do not suggest fields inlined after colon
       }
@@ -84,27 +102,17 @@ export class GxFormat2CompletionService {
           result.push(item);
           return result;
         }
+      } else if (schemaNode.isUnionType) {
+        for (const typeRef of schemaNode.typeRefs) {
+          const typeNode = this.schemaNodeResolver.getSchemaNodeByTypeRef(typeRef);
+          if (typeNode === undefined) continue;
+          result.push(...this.getProposedItems(typeNode, textBuffer, exclude, offset));
+        }
+        return result;
       }
+
       const schemaRecord = this.schemaNodeResolver.getSchemaNodeByTypeRef(schemaNode.typeRef);
-      if (schemaRecord instanceof EnumSchemaNode) {
-        schemaRecord.symbols
-          .filter((v) => v.startsWith(currentWord))
-          .forEach((value) => {
-            if (exclude.has(value)) return;
-            const item: CompletionItem = {
-              label: value,
-              sortText: `_${value}`,
-              kind: CompletionItemKind.EnumMember,
-              documentation: schemaRecord.documentation,
-              insertText: value,
-              textEdit: {
-                range: overwriteRange,
-                newText: value,
-              },
-            };
-            result.push(item);
-          });
-      } else if (schemaRecord instanceof RecordSchemaNode) {
+      if (schemaRecord) {
         return this.getProposedItems(schemaRecord, textBuffer, exclude, offset);
       }
     }

server/gx-workflow-ls-format2/src/services/schemaValidationService.ts

@@ -1,14 +1,15 @@
 import { ASTNodeManager } from "@gxwf/server-common/src/ast/nodeManager";
-import { ASTNode, ObjectASTNode, StringASTNode } from "@gxwf/server-common/src/ast/types";
+import { ASTNode, ObjectASTNode, PropertyASTNode, StringASTNode } from "@gxwf/server-common/src/ast/types";
 import {
   Diagnostic,
   DiagnosticSeverity,
   Range,
   WorkflowDocument,
   WorkflowValidator,
 } from "@gxwf/server-common/src/languageTypes";
+import { isSimpleType } from "@gxwf/server-common/src/utils";
 import { SchemaNode, SchemaNodeResolver } from "../schema";
-import { EnumSchemaNode, IdMapper, RecordSchemaNode } from "../schema/definitions";
+import { EnumSchemaNode, FieldSchemaNode, IdMapper, RecordSchemaNode } from "../schema/definitions";
 
 export class GxFormat2SchemaValidationService implements WorkflowValidator {
   constructor(protected readonly schemaNodeResolver: SchemaNodeResolver) {}
@@ -89,8 +90,10 @@ export class GxFormat2SchemaValidationService implements WorkflowValidator {
           Diagnostic.create(range, `The '${schemaFieldNode.name}' field is required.`, DiagnosticSeverity.Error)
         );
       }
-      if (nodeFound) {
-        if (schemaFieldNode.isPrimitiveType && propertyNode?.valueNode?.type) {
+      if (nodeFound && propertyNode?.valueNode?.type) {
+        const isPropertyTypeSimple = isSimpleType(propertyNode.valueNode.type);
+        // Primitive type validation
+        if (schemaFieldNode.isPrimitiveType && isPropertyTypeSimple) {
           if (!schemaFieldNode.matchesType(propertyNode.valueNode.type)) {
             diagnostics.push(
               Diagnostic.create(
@@ -100,9 +103,30 @@ export class GxFormat2SchemaValidationService implements WorkflowValidator {
               )
             );
           }
+          return;
         }
+
+        // Union type validation
+        if (schemaFieldNode.isUnionType) {
+          if (isPropertyTypeSimple) {
+            const hasMatchingType = this.propetyTypeMatchesAnyPrimitiveRef(schemaFieldNode, propertyNode);
+            if (!hasMatchingType) {
+              diagnostics.push(
+                Diagnostic.create(
+                  range,
+                  `Type mismatch for field '${schemaFieldNode.name}'. Expected '${schemaFieldNode.typeRefs.join(
+                    " | "
+                  )}' but found '${propertyNode.valueNode.type}'.`,
+                  DiagnosticSeverity.Error
+                )
+              );
+            }
+            return;
+          }
+        }
+
         const childSchemaNode = this.schemaNodeResolver.getSchemaNodeByTypeRef(schemaFieldNode.typeRef);
-        if (childSchemaNode && propertyNode.valueNode) {
+        if (childSchemaNode) {
           if (schemaFieldNode.canBeArray) {
             propertyNode.valueNode.children?.forEach((item) => {
               if (item.type === "property" && item.valueNode) {
@@ -117,6 +141,19 @@ export class GxFormat2SchemaValidationService implements WorkflowValidator {
     });
   }
 
+  private propetyTypeMatchesAnyPrimitiveRef(schemaFieldNode: FieldSchemaNode, propertyNode: PropertyASTNode): boolean {
+    let matchesSomeType = false;
+    const possibleTypes = schemaFieldNode.typeRefs;
+    for (const schemaFieldType of possibleTypes) {
+      const isPrimitive = this.schemaNodeResolver.definitions.isPrimitiveType(schemaFieldType);
+      if (isPrimitive && propertyNode.valueNode && schemaFieldNode.matchesType(propertyNode.valueNode.type)) {
+        matchesSomeType = true;
+        break;
+      }
+    }
+    return matchesSomeType;
+  }
+
   private validateNodeTypeDefinition(
     schemaNode: RecordSchemaNode,
     node: ASTNode,

server/gx-workflow-ls-format2/tests/integration/validation.test.ts

@@ -64,22 +64,51 @@ class: GalaxyWorkflow
     expect(diagnostics[2].message).toBe("The 'outputs' field is required.");
   });
 
-  it("should report error for invalid enum value", async () => {
+  it("should report error for invalid input type value", async () => {
     const content = `
 class: GalaxyWorkflow
 inputs:
     the_input:
-        type: unknown
+        type: 5
 outputs:
 steps:
     `;
     const diagnostics = await validateDocument(content);
     expect(diagnostics).toHaveLength(1);
     expect(diagnostics[0].message).toBe(
-      "The value is not a valid 'GalaxyType'. Allowed values are: integer, text, File, data, collection, null, boolean, int, long, float, double, string."
+      "Type mismatch for field 'type'. Expected 'GalaxyType | string' but found 'number'."
     );
   });
 
+  it("should report error for invalid enum value", async () => {
+    const content = `
+class: GalaxyWorkflow
+inputs:
+outputs:
+steps:
+    step:
+      type: unknown
+    `;
+    const diagnostics = await validateDocument(content);
+    expect(diagnostics).toHaveLength(1);
+    expect(diagnostics[0].message).toBe(
+      "The value is not a valid 'WorkflowStepType'. Allowed values are: tool, subworkflow, pause."
+    );
+  });
+
+  it("should not report error for valid enum value", async () => {
+    const content = `
+class: GalaxyWorkflow
+inputs:
+outputs:
+steps:
+    step:
+      type: tool
+    `;
+    const diagnostics = await validateDocument(content);
+    expect(diagnostics).toHaveLength(0);
+  });
+
   it("should not report error for compatible primitive types", async () => {
     const content = `
 class: GalaxyWorkflow
@@ -108,7 +137,9 @@ steps:
     `;
     const diagnostics = await validateDocument(content);
     expect(diagnostics).toHaveLength(1);
-    expect(diagnostics[0].message).toContain("Type mismatch for field 'top'. Expected 'float' but found 'string'.");
+    expect(diagnostics[0].message).toContain(
+      "Type mismatch for field 'top'. Expected 'float | int' but found 'string'."
+    );
   });
 
   it("should not report error for properties with Any type", async () => {
@@ -126,6 +157,35 @@ steps:
     expect(diagnostics).toHaveLength(0);
   });
 
+  it("should not report error when multiple types are allowed (string)", async () => {
+    const content = `
+class: GalaxyWorkflow
+inputs:
+outputs:
+steps:
+    step:
+      out: a string
+    `;
+    const diagnostics = await validateDocument(content);
+    expect(diagnostics).toHaveLength(0);
+  });
+
+  it("should not report error when multiple types are allowed (object)", async () => {
+    const content = `
+class: GalaxyWorkflow
+inputs:
+outputs:
+steps:
+    step:
+      out:
+        add_tags:
+          - tag1
+          - tag2
+    `;
+    const diagnostics = await validateDocument(content);
+    expect(diagnostics).toHaveLength(0);
+  });
+
   describe("Custom Rules", () => {
     let rule: ValidationRule;
 

server/packages/server-common/src/utils.ts

@@ -34,9 +34,26 @@ export function isCompatibleType(expectedType: WorkflowDataType, actualType: str
   return isCompatible;
 }
 
+/**
+ * Check if the type is a valid workflow data type.
+ * @param type The type to check.
+ * @returns True if the type is a valid workflow data type.
+ */
 export function isWorkflowDataType(type?: string): type is WorkflowDataType {
   if (!type) {
     return false;
   }
   return type in workflowDataTypes;
 }
+
+const SIMPLE_TYPES = ["number", "string", "boolean", "null"];
+
+/**
+ * Check if the type is a simple type (i.e. number, string, boolean or null).
+ */
+export function isSimpleType(type?: string): boolean {
+  if (!type) {
+    return false;
+  }
+  return SIMPLE_TYPES.includes(type);
+}