feat(workflow): ParallelWorker and JoinNode (Part 4)

core/src/workflow/nodes/join_node.ts

@@ -0,0 +1,34 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {createEvent, Event} from '../../events/event.js';
+import {BaseNode} from '../base_node.js';
+import {NodeContext} from '../node_context.js';
+
+/**
+ * A fan-in barrier node: it waits for ALL of its predecessors to complete, then
+ * emits the aggregated inputs (a map of predecessor name → output) as its
+ * output.
+ *
+ * Ported from `google/adk-python` `workflow/_join_node.py`.
+ */
+export class JoinNode extends BaseNode {
+  override get requiresAllPredecessors(): boolean {
+    return true;
+  }
+
+  protected async *runImpl(
+    ctx: NodeContext,
+    input: unknown,
+  ): AsyncGenerator<Event, void, void> {
+    yield createEvent({
+      author: this.name,
+      invocationId: ctx.invocationId,
+      branch: ctx.branch,
+      output: input,
+    });
+  }
+}

core/src/workflow/nodes/parallel_worker.ts

@@ -0,0 +1,113 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {BaseNode} from '../base_node.js';
+import {NodeContext} from '../node_context.js';
+import {RetryConfig} from '../retry_config.js';
+import {registerParallelWorkerFactory} from '../utils/workflow_graph_utils.js';
+
+/** Options for a {@link ParallelWorker}. */
+export interface ParallelWorkerConfig {
+  /** Maximum number of items processed concurrently. `undefined` = unlimited. */
+  maxParallelWorkers?: number;
+  retryConfig?: RetryConfig;
+  timeout?: number;
+}
+
+/**
+ * A node that runs a wrapped node in parallel for each item of a list input,
+ * preserving order, bounded by `maxParallelWorkers`, cancelling on first error.
+ *
+ * Ported from `google/adk-python` `workflow/_parallel_worker.py`. A non-list
+ * input is treated as a single-element list. Each item runs via
+ * `ctx.runNode(inner, item, {useSubBranch: true})`; the node's output is the
+ * ordered list of the children's outputs.
+ */
+export class ParallelWorker extends BaseNode {
+  readonly maxParallelWorkers?: number;
+  private readonly inner: BaseNode;
+
+  constructor(inner: BaseNode, config: ParallelWorkerConfig = {}) {
+    super({
+      name: inner.name,
+      rerunOnResume: true,
+      retryConfig: config.retryConfig,
+      timeout: config.timeout,
+    });
+    if (
+      config.maxParallelWorkers !== undefined &&
+      config.maxParallelWorkers < 1
+    ) {
+      throw new Error('maxParallelWorkers must be greater than or equal to 1.');
+    }
+    this.inner = inner;
+    this.maxParallelWorkers = config.maxParallelWorkers;
+  }
+
+  protected async *runImpl(
+    ctx: NodeContext,
+    input: unknown,
+  ): AsyncGenerator<unknown, void, void> {
+    const items = Array.isArray(input) ? input : [input];
+    if (items.length === 0) {
+      yield [];
+      return;
+    }
+
+    const results = new Array<unknown>(items.length);
+    const poolSize = Math.min(
+      this.maxParallelWorkers ?? items.length,
+      items.length,
+    );
+
+    let nextIndex = 0;
+    let firstError: unknown;
+
+    const worker = async (): Promise<void> => {
+      for (;;) {
+        if (firstError !== undefined) {
+          return;
+        }
+        const i = nextIndex++;
+        if (i >= items.length) {
+          return;
+        }
+        try {
+          // Key each child run by its item index (not call order) so the
+          // run id -> item mapping is deterministic. On resume this lets each
+          // item fast-forward from its own cached run rather than being matched
+          // to a differently-ordered run id.
+          const child = await ctx.runNode(this.inner, items[i], {
+            useSubBranch: true,
+            runId: String(i),
+          });
+          results[i] = child.output;
+        } catch (err) {
+          if (firstError === undefined) {
+            firstError = err;
+          }
+          return;
+        }
+      }
+    };
+
+    await Promise.all(Array.from({length: poolSize}, () => worker()));
+
+    if (firstError !== undefined) {
+      throw firstError;
+    }
+    yield results;
+  }
+}
+
+/**
+ * Registers the factory the engine uses to wrap a built node in a
+ * {@link ParallelWorker} when `buildNode(..., {parallelWorker: true})` is
+ * requested — keeping the engine core free of a static import of this module.
+ */
+registerParallelWorkerFactory(
+  (inner, options) => new ParallelWorker(inner, options),
+);

core/test/workflow/parallel_worker_test.ts

@@ -0,0 +1,95 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {describe, expect, it} from 'vitest';
+import {FunctionNode} from '../../src/workflow/nodes/function_node.js';
+import {JoinNode} from '../../src/workflow/nodes/join_node.js';
+import {ParallelWorker} from '../../src/workflow/nodes/parallel_worker.js';
+import {buildNode} from '../../src/workflow/utils/workflow_graph_utils.js';
+import {driveNode} from './test_helpers.js';
+
+describe('ParallelWorker', () => {
+  it('maps a list input through the inner node, preserving order', async () => {
+    const inner = new FunctionNode('double', (_c, n: number) => n * 2);
+    const {output} = await driveNode(new ParallelWorker(inner), [1, 2, 3, 4]);
+    expect(output).toEqual([2, 4, 6, 8]);
+  });
+
+  it('treats a non-list input as a single-element list', async () => {
+    const inner = new FunctionNode('double', (_c, n: number) => n * 2);
+    const {output} = await driveNode(new ParallelWorker(inner), 5);
+    expect(output).toEqual([10]);
+  });
+
+  it('yields an empty list for an empty input', async () => {
+    const inner = new FunctionNode('id', (_c, x) => x);
+    const {output} = await driveNode(new ParallelWorker(inner), []);
+    expect(output).toEqual([]);
+  });
+
+  it('bounds concurrency by maxParallelWorkers', async () => {
+    let active = 0;
+    let peak = 0;
+    const inner = new FunctionNode('track', async (_c, n: number) => {
+      active++;
+      peak = Math.max(peak, active);
+      await new Promise((r) => setTimeout(r, 5));
+      active--;
+      return n;
+    });
+    const {output} = await driveNode(
+      new ParallelWorker(inner, {maxParallelWorkers: 2}),
+      [1, 2, 3, 4, 5],
+    );
+    expect(output).toEqual([1, 2, 3, 4, 5]);
+    expect(peak).toBeLessThanOrEqual(2);
+  });
+
+  it('rejects maxParallelWorkers < 1', () => {
+    const inner = new FunctionNode('x', (_c, v) => v);
+    expect(() => new ParallelWorker(inner, {maxParallelWorkers: 0})).toThrow(
+      /greater than or equal to 1/,
+    );
+  });
+
+  it('propagates the first error from a failing item', async () => {
+    const inner = new FunctionNode('boom', (_c, n: number) => {
+      if (n === 3) {
+        throw new Error('boom at 3');
+      }
+      return n;
+    });
+    await expect(
+      driveNode(new ParallelWorker(inner), [1, 2, 3, 4]),
+    ).rejects.toThrow('boom at 3');
+  });
+});
+
+describe('ParallelWorker registry factory', () => {
+  it('buildNode wraps the built node when parallelWorker is requested', () => {
+    const node = buildNode((_c: unknown, n: number) => n, {
+      name: 'w',
+      parallelWorker: true,
+    });
+    expect(node).toBeInstanceOf(ParallelWorker);
+  });
+
+  it('rejects maxParallelWorkers without parallelWorker', () => {
+    expect(() =>
+      buildNode(() => {}, {name: 'x', maxParallelWorkers: 2}),
+    ).toThrow(/maxParallelWorkers can only be set/);
+  });
+});
+
+describe('JoinNode', () => {
+  it('emits its aggregated input as output and requires all predecessors', async () => {
+    const join = new JoinNode({name: 'join'});
+    const aggregated = {a: 1, b: 2};
+    const {output} = await driveNode(join, aggregated);
+    expect(output).toEqual(aggregated);
+    expect(join.requiresAllPredecessors).toBe(true);
+  });
+});