[refactor] Enhance ConditionalTask and WhileTask configuration handling

- Updated ConditionalTask to include a new optional `conditionConfig` property in its configuration schema, allowing for more flexible branch conditions.
- Refactored WhileTask to remove reliance on `extras` for condition fields, integrating them directly into the task configuration for improved clarity and usability.
- Adjusted related types and serialization structures to accommodate the new configuration format.
- Updated tests to validate the new configuration handling for both ConditionalTask and WhileTask, ensuring consistent behavior across task executions.

Author: sroussey

packages/task-graph/src/task/ConditionalTask.ts

@@ -98,15 +98,18 @@ export const conditionalTaskConfigSchema = {
     branches: { type: "array", items: {} },
     defaultBranch: { type: "string" },
     exclusive: { type: "boolean" },
+    conditionConfig: { type: "object", additionalProperties: true },
   },
   additionalProperties: false,
 } as const satisfies DataPortSchema;
 
 export type ConditionalTaskConfig = TaskConfig & {
   /** Branches may contain ConditionFn functions — not JSON-schema-representable */
-  readonly branches: BranchConfig<any>[];
+  readonly branches?: BranchConfig<any>[];
   readonly defaultBranch?: string;
   readonly exclusive?: boolean;
+  /** Serializable UI condition configuration used to build branches at runtime. */
+  readonly conditionConfig?: UIConditionConfig;
 };
 
 // ============================================================================
@@ -283,10 +286,10 @@ export class ConditionalTask<
       };
     }
 
-    // Try to find serialized conditionConfig from input or extras
+    // Try to find serialized conditionConfig from input or config
     const conditionConfig =
       ((input as Record<string, unknown>).conditionConfig as UIConditionConfig | undefined) ??
-      (this.config.extras?.conditionConfig as UIConditionConfig | undefined);
+      this.config.conditionConfig;
 
     if (conditionConfig) {
       return {

packages/task-graph/src/task/TaskJSON.ts

@@ -9,8 +9,11 @@ import { TaskGraph } from "../task-graph/TaskGraph";
 import { CompoundMergeStrategy } from "../task-graph/TaskGraphRunner";
 import { TaskConfigurationError, TaskJSONError } from "../task/TaskError";
 import { TaskRegistry } from "../task/TaskRegistry";
+import { ConditionalTaskConfig } from "./ConditionalTask";
 import { GraphAsTask } from "./GraphAsTask";
+import { IteratorTaskConfig } from "./IteratorTask";
 import { TaskConfig, TaskInput } from "./TaskTypes";
+import { WhileTaskConfig } from "./WhileTask";
 
 // ========================================================================
 // JSON Serialization Types
@@ -20,6 +23,11 @@ import { TaskConfig, TaskInput } from "./TaskTypes";
  * This structure defines how tasks should be configured in JSON format.
  */
 
+export type JsonTaskConfig = Omit<
+  TaskConfig & WhileTaskConfig & IteratorTaskConfig & ConditionalTaskConfig,
+  "id"
+>;
+
 export type JsonTaskItem = {
   /** Unique identifier for the task */
   id: unknown;
@@ -28,7 +36,7 @@ export type JsonTaskItem = {
   type: string;
 
   /** Optional configuration for the task */
-  config?: Omit<TaskConfig, "id"> & Record<string, unknown>;
+  config?: JsonTaskConfig;
 
   /** Default input values for the task */
   defaults?: TaskInput;
@@ -61,7 +69,7 @@ export type TaskGraphItemJson = {
   id: unknown;
   type: string;
   defaults?: TaskInput;
-  config?: Omit<TaskConfig, "id"> & Record<string, unknown>;
+  config?: JsonTaskConfig;
   subgraph?: TaskGraphJson;
   merge?: CompoundMergeStrategy;
 };

packages/task-graph/src/task/TaskTypes.ts

@@ -94,7 +94,10 @@ export const TaskConfigSchema = {
     cacheable: { type: "boolean" },
     inputSchema: { type: "object", properties: {}, additionalProperties: true },
     outputSchema: { type: "object", properties: {}, additionalProperties: true },
-    extras: { type: "object", additionalProperties: true },
+    extras: {
+      type: "object",
+      additionalProperties: true,
+    },
   },
   additionalProperties: false,
 } as const satisfies DataPortSchema;

packages/task-graph/src/task/WhileTask.ts

@@ -49,6 +49,10 @@ export const whileTaskConfigSchema = {
     condition: {},
     maxIterations: { type: "integer", minimum: 1 },
     chainIterations: { type: "boolean" },
+    conditionField: { type: "string" },
+    conditionOperator: { type: "string" },
+    conditionValue: { type: "string" },
+    iterationInputConfig: { type: "object", additionalProperties: true },
   },
   additionalProperties: false,
 } as const satisfies DataPortSchema;
@@ -76,6 +80,18 @@ export type WhileTaskConfig<Output extends TaskOutput = TaskOutput> = GraphAsTas
    * @default true
    */
   readonly chainIterations?: boolean;
+
+  /** Output field to evaluate for the loop condition. */
+  readonly conditionField?: string;
+
+  /** Comparison operator for the loop condition. */
+  readonly conditionOperator?: string;
+
+  /** Value to compare against for the loop condition. */
+  readonly conditionValue?: string;
+
+  /** Per-property iteration input configuration (scalar/array/flexible). */
+  readonly iterationInputConfig?: Record<string, { mode: string; baseSchema?: unknown }>;
 };
 
 /**
@@ -184,22 +200,16 @@ export class WhileTask<
 
   /**
    * Gets the maximum iterations limit.
-   * Falls back to extras.whileConfig.maxIterations for JSON-deserialized tasks.
    */
   public get maxIterations(): number {
-    if (this.config.maxIterations !== undefined) return this.config.maxIterations;
-    const wc = this.config.extras?.whileConfig as { maxIterations?: number } | undefined;
-    return wc?.maxIterations ?? 100;
+    return this.config.maxIterations ?? 100;
   }
 
   /**
    * Whether to chain iteration outputs to inputs.
-   * Falls back to extras.whileConfig.chainIterations for JSON-deserialized tasks.
    */
   public get chainIterations(): boolean {
-    if (this.config.chainIterations !== undefined) return this.config.chainIterations;
-    const wc = this.config.extras?.whileConfig as { chainIterations?: boolean } | undefined;
-    return wc?.chainIterations ?? true;
+    return this.config.chainIterations ?? true;
   }
 
   /**
@@ -217,20 +227,15 @@ export class WhileTask<
    * Execute the while loop.
    */
   /**
-   * Builds a condition function from the serialized whileConfig in extras
-   * when no condition function is directly provided in config.
+   * Builds a condition function from the serialized condition fields in config.
    */
-  private buildConditionFromExtras(): WhileConditionFn<Output> | undefined {
-    const wc = this.config.extras?.whileConfig as
-      | { conditionField?: string; conditionOperator?: string; conditionValue?: string }
-      | undefined;
+  private buildConditionFromConfig(): WhileConditionFn<Output> | undefined {
+    const { conditionOperator, conditionField, conditionValue } = this.config;
 
-    if (!wc?.conditionOperator) {
+    if (!conditionOperator) {
       return undefined;
     }
 
-    const { conditionField, conditionOperator, conditionValue } = wc;
-
     return (output: Output) => {
       const fieldValue = conditionField
         ? getNestedValue(output as Record<string, unknown>, conditionField)
@@ -251,16 +256,12 @@ export class WhileTask<
     iteratedValues: Record<string, unknown[]>;
     iterationCount: number;
   } | null {
-    const wc = this.config.extras?.whileConfig as
-      | { iterationInputConfig?: Record<string, { mode: string; baseSchema?: unknown }> }
-      | undefined;
-
-    if (!wc?.iterationInputConfig) {
+    if (!this.config.iterationInputConfig) {
       return null;
     }
 
     const inputData = input as Record<string, unknown>;
-    const config = wc.iterationInputConfig;
+    const config = this.config.iterationInputConfig!;
 
     const arrayPorts: string[] = [];
     const scalarPorts: string[] = [];
@@ -350,7 +351,7 @@ export class WhileTask<
     }
 
     // Use provided condition or auto-build from serialized whileConfig
-    const condition = this.condition ?? this.buildConditionFromExtras();
+    const condition = this.condition ?? this.buildConditionFromConfig();
 
     if (!condition) {
       throw new TaskConfigurationError(`${this.type}: No condition function provided`);
@@ -426,15 +427,12 @@ export class WhileTask<
    * This provides streaming output for the final result while still
    * supporting iteration chaining.
    */
-  async *executeStream(
-    input: Input,
-    context: IExecuteContext
-  ): AsyncIterable<StreamEvent<Output>> {
+  async *executeStream(input: Input, context: IExecuteContext): AsyncIterable<StreamEvent<Output>> {
     if (!this.hasChildren()) {
       throw new TaskConfigurationError(`${this.type}: No subgraph set for while loop`);
     }
 
-    const condition = this.condition ?? this.buildConditionFromExtras();
+    const condition = this.condition ?? this.buildConditionFromConfig();
     if (!condition) {
       throw new TaskConfigurationError(`${this.type}: No condition function provided`);
     }
@@ -550,19 +548,15 @@ export class WhileTask<
     const baseSchema = super.inputSchema();
     if (typeof baseSchema === "boolean") return baseSchema;
 
-    const wc = this.config.extras?.whileConfig as
-      | { iterationInputConfig?: Record<string, { mode: string; baseSchema?: DataPortSchema }> }
-      | undefined;
-
-    if (!wc?.iterationInputConfig) {
+    if (!this.config.iterationInputConfig) {
       return baseSchema;
     }
 
     // Wrap array-mode ports in anyOf (scalar | array) schemas.
     // Using anyOf instead of plain type:"array" to avoid addInput's array-merge behavior
     // which would prepend an undefined element when runInputData starts empty.
     const properties = { ...(baseSchema.properties || {}) } as Record<string, DataPortSchema>;
-    for (const [key, propConfig] of Object.entries(wc.iterationInputConfig)) {
+    for (const [key, propConfig] of Object.entries(this.config.iterationInputConfig)) {
       if (propConfig.mode === "array" && properties[key]) {
         const scalarSchema = properties[key] as DataPortSchema;
         properties[key] = {

packages/test/src/test/task/ConditionalTaskCondition.test.ts

@@ -343,21 +343,18 @@ describe("ConditionalTask with serialized conditionConfig", () => {
     });
   });
 
-  describe("conditionConfig from extras", () => {
-    it("should use conditionConfig from extras when input has none", async () => {
-      const conditionConfig: UIConditionConfig = {
-        branches: [
-          { id: "high", field: "value", operator: "greater_than", value: "50" },
-          { id: "low", field: "value", operator: "less_or_equal", value: "50" },
-        ],
-        exclusive: true,
-      };
-
+  describe("conditionConfig from config", () => {
+    it("should use conditionConfig from config when input has none", async () => {
       const task = new ConditionalTask(
         {},
         {
-          branches: [],
-          extras: { conditionConfig },
+          conditionConfig: {
+            branches: [
+              { id: "high", field: "value", operator: "greater_than", value: "50" },
+              { id: "low", field: "value", operator: "less_or_equal", value: "50" },
+            ],
+            exclusive: true,
+          },
         }
       );
 

packages/test/src/test/task/WhileTaskCondition.test.ts

@@ -20,13 +20,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 20,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "less_than",
-              conditionValue: "0.9",
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "less_than",
+          conditionValue: "0.9",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -46,13 +42,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 20,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "not_equals",
-              conditionValue: "1",
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "not_equals",
+          conditionValue: "1",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -70,13 +62,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 20,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "equals",
-              conditionValue: "0.2",
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "equals",
+          conditionValue: "0.2",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -96,13 +84,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 3,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "less_than",
-              conditionValue: "999", // always true
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "less_than",
+          conditionValue: "999",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -122,13 +106,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 3,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "less_than",
-              conditionValue: "999",
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "less_than",
+          conditionValue: "999",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -148,13 +128,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 3,
           chainIterations: false,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "less_than",
-              conditionValue: "999",
-            },
-          },
+          conditionField: "quality",
+          conditionOperator: "less_than",
+          conditionValue: "999",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -177,13 +153,9 @@ describe("WhileTask with serialized whileConfig", () => {
         .while({
           maxIterations: 3,
           chainIterations: true,
-          extras: {
-            whileConfig: {
-              conditionField: "",
-              conditionOperator: "is_not_empty",
-              conditionValue: "",
-            },
-          },
+          conditionField: "",
+          conditionOperator: "is_not_empty",
+          conditionValue: "",
         })
         .addTask(RefineTask)
         .endWhile();
@@ -203,13 +175,6 @@ describe("WhileTask with serialized whileConfig", () => {
           condition: (output: { quality: number }, iteration: number) =>
             output.quality < 0.5 && iteration < 10,
           maxIterations: 20,
-          extras: {
-            whileConfig: {
-              conditionField: "quality",
-              conditionOperator: "less_than",
-              conditionValue: "999", // would loop forever without maxIterations
-            },
-          },
         })
         .addTask(RefineTask)
         .endWhile();