From 528a99d4d79c73f5f654d7f0f7777b41e32fd46c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Thu, 29 Oct 2026 06:21:19 -0700 Subject: [PATCH] feat: run a node once per item of a list input PiperOrigin-RevId: 1009638966 --- .../google/adk/workflow/ParallelWorkers.java | 228 ++++++ .../adk/workflow/ParallelWorkersTest.java | 746 ++++++++++++++++++ 2 files changed, 974 insertions(+) create mode 100644 core/src/main/java/com/google/adk/workflow/ParallelWorkers.java create mode 100644 core/src/test/java/com/google/adk/workflow/ParallelWorkersTest.java diff --git a/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java b/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java new file mode 100644 index 000000000..129d4f5c6 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java @@ -0,0 +1,228 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.workflow; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectReader; +import com.google.adk.JsonBaseModel; +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReferenceArray; +import org.jspecify.annotations.Nullable; + +/** + * Static factories for workflow nodes that run a node once per item of a list input. Each item runs + * as a child of the returned node on its own branch, and the node outputs the items' results as a + * list in input order. A failed item cancels the others and fails the node, while an item that asks + * for input stops queued items from starting. + */ +@Experimental +public final class ParallelWorkers { + + /** + * Returns a node that runs {@code node} on every item at once. Otherwise it behaves as {@link + * #parallelNode(Node, int)} describes. + * + * @throws WorkflowConfigurationException if {@code node} is {@link Start} + */ + public static Node parallelNode(Node node) { + return parallelNode(node, Integer.MAX_VALUE); + } + + /** + * Returns a node that runs {@code node} on each item of a list input, at most {@code + * maxParallelWorkers} at a time and starting them in input order; the returned node takes {@code + * node}'s name, description, state schema and {@link Node#requiresAllPredecessors()}, and leaves + * its config and input and output schemas to each item's run. A text-only {@link Content}, such + * as a user message, is read as JSON, and a JSON list fans out. Its text reaches the items as a + * {@code String} if it is not JSON, or if it is a non-list value and the item's input schema + * takes a string; any other non-list input runs as one item. + * + * @throws WorkflowConfigurationException if {@code node} is {@link Start} or {@code + * maxParallelWorkers} is less than 1 + */ + public static Node parallelNode(Node node, int maxParallelWorkers) { + if (node instanceof Start) { + throw new WorkflowConfigurationException("parallelNode cannot wrap Start."); + } + if (maxParallelWorkers < 1) { + throw new WorkflowConfigurationException( + String.format("maxParallelWorkers must be at least 1, but was %s.", maxParallelWorkers)); + } + return new ParallelWorkerNode(node, maxParallelWorkers); + } + + private static final class ParallelWorkerNode extends BaseNode { + + private static final ObjectReader JSON_READER = + JsonBaseModel.getMapper() + .readerFor(Object.class) + .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + // No run ID: items subscribe in input order, so each gets the same generated ID on resume. + private static final RunNodeConfig ITEM_CONFIG = + RunNodeConfig.builder().useSubBranch(true).build(); + + private final Node worker; + private final int maxParallelWorkers; + + ParallelWorkerNode(Node worker, int maxParallelWorkers) { + super( + worker.name(), + worker.description(), + /* rerunOnResume= */ true, + /* waitForOutput= */ false, + NodeConfig.builder().build(), + /* inputSchema= */ null, + /* outputSchema= */ null, + worker.stateSchema().orElse(null)); + this.worker = worker; + this.maxParallelWorkers = maxParallelWorkers; + } + + @Override + public boolean requiresAllPredecessors() { + return worker.requiresAllPredecessors(); + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.defer( + () -> { + Object input = nodeInput instanceof Content content ? unwrap(content) : nodeInput; + List items = input instanceof List list ? list : Collections.singletonList(input); + return new FanOut(context, items).run(); + }); + } + + /** + * Returns the JSON value of text-only {@code content}, or its text if that is not JSON or is a + * non-list value for an item that takes a string; any other content is returned as it is. + */ + private @Nullable Object unwrap(Content content) { + List parts = content.parts().orElse(ImmutableList.of()); + if (parts.isEmpty() || parts.stream().anyMatch(part -> part.text().isEmpty())) { + return content; + } + String text = content.text(); + Object json; + try { + json = JSON_READER.readValue(text); + } catch (JsonProcessingException e) { + return text; + } + if (json instanceof List list) { + return list; + } + return worker.inputSchema().map(ParallelWorkerNode::acceptsString).orElse(false) + ? text + : json; + } + + private static boolean acceptsString(Schema schema) { + return schema.type().map(type -> type.knownEnum() == Type.Known.STRING).orElse(false) + || schema.anyOf().orElse(ImmutableList.of()).stream() + .anyMatch(ParallelWorkerNode::acceptsString); + } + + private final class FanOut { + private final Context context; + private final List items; + // Written by the item runs, which may finish on different threads. + private final AtomicReferenceArray outputs; + private final AtomicReferenceArray failures; + private final AtomicBoolean interrupted = new AtomicBoolean(); + + FanOut(Context context, List items) { + this.context = context; + this.items = items; + this.outputs = new AtomicReferenceArray<>(items.size()); + this.failures = new AtomicReferenceArray<>(items.size()); + } + + Flowable run() { + return Flowable.range(0, items.size()) + // An interrupt stops only the items that had to wait for a slot. + .takeWhile(index -> index < maxParallelWorkers || !interrupted.get()) + .flatMapSingle(this::runItem, /* delayErrors= */ false, maxParallelWorkers) + // A failed item stops the rest; the typed parameter rules out takeUntil(Publisher). + .takeUntil((Boolean failed) -> failed) + .ignoreElements() + .andThen(Flowable.defer(this::finish)); + } + + /** Runs one item, or replays its result on resume, and emits whether it failed. */ + private Single runItem(int index) { + return context + .runNode(worker, items.get(index), ITEM_CONFIG) + .doOnSuccess(output -> outputs.set(index, output)) + .ignoreElement() + .toSingleDefault(false) + .onErrorResumeNext(error -> settle(index, error)); + } + + /** + * Records a failure so the lowest index wins a race, and keeps an interrupt from ending the + * rest. + */ + private Single settle(int index, Throwable error) { + if (error instanceof DynamicNodeFailedException failure) { + failures.set(index, failure); + return Single.just(true); + } + if (error instanceof NodeInterruptedException) { + // The runtime has already added the item's interrupt IDs to this node. + interrupted.set(true); + return Single.just(false); + } + return Single.error(error); + } + + /** + * Fails with the lowest-index failure, or else waits with the items that asked for input, or + * else outputs the items' results. Each failure's error event stays at its item. + */ + private Flowable finish() { + List<@Nullable Object> results = new ArrayList<>(items.size()); + for (int i = 0; i < items.size(); i++) { + DynamicNodeFailedException failure = failures.get(i); + if (failure != null) { + return Flowable.error(failure); + } + results.add(outputs.get(i)); + } + if (interrupted.get()) { + return Flowable.error(new NodeInterruptedException()); + } + return Flowable.just(Collections.unmodifiableList(results)); + } + } + } + + private ParallelWorkers() {} +} diff --git a/core/src/test/java/com/google/adk/workflow/ParallelWorkersTest.java b/core/src/test/java/com/google/adk/workflow/ParallelWorkersTest.java new file mode 100644 index 000000000..cbd280ba3 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/ParallelWorkersTest.java @@ -0,0 +1,746 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.workflow; + +import static com.google.adk.testing.WorkflowTestNodes.emitter; +import static com.google.adk.testing.WorkflowTestNodes.workflow; +import static com.google.adk.workflow.WorkflowResumeTesting.START; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.interrupt; +import static com.google.adk.workflow.WorkflowResumeTesting.outputEventAt; +import static com.google.adk.workflow.WorkflowResumeTesting.pathOf; +import static com.google.adk.workflow.WorkflowResumeTesting.pathsOf; +import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.Context; +import com.google.adk.events.Event; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.observers.TestObserver; +import io.reactivex.rxjava3.schedulers.Schedulers; +import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.time.Duration; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Function; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class ParallelWorkersTest { + + private static final Schema INPUT_SCHEMA = Schema.builder().type("ARRAY").build(); + private static final Schema OUTPUT_SCHEMA = Schema.builder().type("STRING").build(); + private static final Schema STATE_SCHEMA = Schema.builder().type("OBJECT").build(); + + /** A node with a description, config, and input, output and state schemas of its own. */ + private static final Node DESCRIBED_WORKER = + new Node() { + @Override + public String name() { + return "worker"; + } + + @Override + public String description() { + return "A worker with its own settings."; + } + + @Override + public NodeConfig config() { + return NodeConfig.builder().timeout(Duration.ofSeconds(1)).build(); + } + + @Override + public Optional inputSchema() { + return Optional.of(INPUT_SCHEMA); + } + + @Override + public Optional outputSchema() { + return Optional.of(OUTPUT_SCHEMA); + } + + @Override + public Optional stateSchema() { + return Optional.of(STATE_SCHEMA); + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.empty(); + } + }; + + @Test + public void parallelNode_takesTheNodesNameDescriptionAndStateSchema() { + Node node = ParallelWorkers.parallelNode(DESCRIBED_WORKER); + + assertThat(node.name()).isEqualTo("worker"); + assertThat(node.description()).isEqualTo("A worker with its own settings."); + assertThat(node.stateSchema()).hasValue(STATE_SCHEMA); + } + + @Test + public void parallelNode_rerunsOnResumeWithoutWaitingForOutput() { + Node node = ParallelWorkers.parallelNode(DESCRIBED_WORKER); + + assertThat(node.rerunOnResume()).isTrue(); + assertThat(node.waitForOutput()).isFalse(); + } + + @Test + public void parallelNode_leavesConfigAndSchemasToEachItem() { + Node node = ParallelWorkers.parallelNode(DESCRIBED_WORKER); + + assertThat(node.config()).isEqualTo(NodeConfig.builder().build()); + assertThat(node.inputSchema()).isEmpty(); + assertThat(node.outputSchema()).isEmpty(); + } + + @Test + public void parallelNode_takesRequiresAllPredecessorsFromTheNode() { + Node join = JoinNode.builder().name("gather").build(); + + assertThat(ParallelWorkers.parallelNode(join).requiresAllPredecessors()).isTrue(); + assertThat(ParallelWorkers.parallelNode(doubler()).requiresAllPredecessors()).isFalse(); + } + + @Test + public void parallelNode_start_isRejected() { + WorkflowConfigurationException thrown = + assertThrows( + WorkflowConfigurationException.class, + () -> ParallelWorkers.parallelNode(Start.INSTANCE)); + + assertThat(thrown).hasMessageThat().contains("cannot wrap Start"); + } + + @Test + public void parallelNode_capBelowOne_isRejected() { + Node worker = doubler(); + + WorkflowConfigurationException thrown = + assertThrows( + WorkflowConfigurationException.class, + () -> ParallelWorkers.parallelNode(worker, /* maxParallelWorkers= */ 0)); + + assertThat(thrown).hasMessageThat().contains("maxParallelWorkers must be at least 1"); + } + + @Test + public void runNode_listInput_runsEachItemOnItsOwnBranchAndOutputsInInputOrder() { + Run run = run(ParallelWorkers.parallelNode(doubler()), ImmutableList.of(1, 2, 3)); + + assertThat(run.output()).isEqualTo(ImmutableList.of(2, 4, 6)); + ImmutableList itemEvents = run.eventsUnder("worker@1/"); + assertThat(pathsOf(itemEvents)) + .containsExactly("worker@1/worker@1", "worker@1/worker@2", "worker@1/worker@3") + .inOrder(); + assertThat( + itemEvents.stream().map(event -> event.branch().orElse("")).collect(toImmutableList())) + .containsExactly("worker@1", "worker@2", "worker@3") + .inOrder(); + } + + @Test + public void runNode_laterItemFinishesFirst_keepsInputOrder() { + CountDownLatch secondDone = new CountDownLatch(1); + Node worker = + node( + item -> + item.equals(1) + ? Flowable.fromCallable( + () -> { + awaitRelease(secondDone); + return "res-1"; + }) + .subscribeOn(Schedulers.io()) + : Flowable.just("res-2").doOnComplete(secondDone::countDown)); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2)); + + assertThat(pathsOf(run.eventsUnder("worker@1/"))) + .containsExactly("worker@1/worker@2", "worker@1/worker@1") + .inOrder(); + assertThat(run.output()).isEqualTo(ImmutableList.of("res-1", "res-2")); + } + + @Test + public void runNode_valueThatIsNotAList_runsAsOneItem() { + Run run = run(ParallelWorkers.parallelNode(doubler()), 5); + + assertThat(run.output()).isEqualTo(ImmutableList.of(10)); + } + + @Test + public void runNode_nullInput_runsAsOneItem() { + Node worker = node(item -> Flowable.just("got " + item)); + + Run run = run(ParallelWorkers.parallelNode(worker), null); + + assertThat(run.output()).isEqualTo(ImmutableList.of("got null")); + } + + @Test + public void runNode_itemWithoutOutput_leavesNullInItsPlace() { + Node worker = node(item -> item.equals(2) ? Flowable.empty() : Flowable.just(item)); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2, 3)); + + assertThat(run.output()).isEqualTo(Arrays.asList(1, null, 3)); + } + + @Test + public void runNode_emptyList_runsNothingAndOutputsAnEmptyList() { + Run run = run(ParallelWorkers.parallelNode(doubler()), ImmutableList.of()); + + assertThat(run.output()).isEqualTo(ImmutableList.of()); + assertThat(run.eventsUnder("worker@1/")).isEmpty(); + } + + @Test + public void runNode_maxParallelWorkers_startsNoMoreItemsUntilOneFinishes() { + AtomicInteger started = new AtomicInteger(); + CompletableSubject release = CompletableSubject.create(); + Node worker = + node( + item -> { + started.incrementAndGet(); + return release.andThen(Flowable.just(item)); + }); + TestObserver observer = + start( + ParallelWorkers.parallelNode(worker, /* maxParallelWorkers= */ 2), + ImmutableList.of(1, 2, 3, 4, 5, 6)) + .test(); + int startedBeforeRelease = started.get(); + + release.onComplete(); + + assertThat(startedBeforeRelease).isEqualTo(2); + assertThat(observer.awaitDone(10, SECONDS).values().get(0).requireNodeState().output()) + .hasValue(ImmutableList.of(1, 2, 3, 4, 5, 6)); + } + + @Test + public void runNode_noMaxParallelWorkers_startsEveryItemAtOnce() { + AtomicInteger started = new AtomicInteger(); + CompletableSubject release = CompletableSubject.create(); + Node worker = + node( + item -> { + started.incrementAndGet(); + return release.andThen(Flowable.just(item)); + }); + TestObserver observer = + start(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2, 3)).test(); + int startedBeforeRelease = started.get(); + + release.onComplete(); + + assertThat(startedBeforeRelease).isEqualTo(3); + assertThat(observer.awaitDone(10, SECONDS).values().get(0).requireNodeState().output()) + .hasValue(ImmutableList.of(1, 2, 3)); + } + + @Test + public void runNode_itemConfigWithRetries_retriesOnlyTheFailedItem() { + Map attemptsByItem = new ConcurrentHashMap<>(); + NodeConfig retryOnce = + NodeConfig.builder() + .retryConfig(RetryConfig.builder().maxAttempts(2).initialDelay(Duration.ZERO).build()) + .build(); + Node worker = + node( + retryOnce, + item -> + attemptsByItem.merge(item, 1, Integer::sum) == 1 && item.equals(2) + ? Flowable.error(new IllegalStateException("first attempt failed")) + : Flowable.just(item)); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2)); + + assertThat(run.output()).isEqualTo(ImmutableList.of(1, 2)); + assertThat(attemptsByItem).containsExactly(1, 1, 2, 2); + } + + @Test + public void runNode_itemConfigWithTimeout_failsTheItemThatRunsTooLong() { + Node worker = + node( + NodeConfig.builder().timeout(Duration.ofMillis(100)).build(), item -> Flowable.never()); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of("slow")); + + assertThat(run.failure().orElseThrow()).isInstanceOf(TimeoutException.class); + assertThat(run.context().requireNodeState().failureNodePath()).hasValue("worker@1/worker@1"); + } + + @Test + public void runNode_appliedTwice_fansOutOverNestedLists() { + Node node = ParallelWorkers.parallelNode(ParallelWorkers.parallelNode(doubler())); + + Run run = run(node, ImmutableList.of(ImmutableList.of(1, 2), ImmutableList.of(3))); + + assertThat(run.output()) + .isEqualTo(ImmutableList.of(ImmutableList.of(2, 4), ImmutableList.of(6))); + } + + @Test + public void runNode_itemFails_failsTheNodeWithOneErrorEventAtTheItem() { + Run run = run(ParallelWorkers.parallelNode(failsOn(2)), ImmutableList.of(1, 2, 3)); + + assertThat(run.failure().orElseThrow()).hasMessageThat().isEqualTo("bad item 2"); + assertThat(run.context().requireNodeState().failureNodePath()).hasValue("worker@1/worker@2"); + assertThat(run.output()).isNull(); + assertThat( + pathsOf( + run.events().stream() + .filter(event -> event.errorCode().isPresent()) + .collect(toImmutableList()))) + .containsExactly("worker@1/worker@2"); + } + + @Test + public void runNode_itemFailsUnderCap_neverStartsQueuedItems() { + Run run = + run( + ParallelWorkers.parallelNode(failsOn(1), /* maxParallelWorkers= */ 1), + ImmutableList.of(1, 2, 3)); + + assertThat(run.failure()).isPresent(); + assertThat(ImmutableSet.copyOf(pathsOf(run.eventsUnder("worker@1/")))) + .containsExactly("worker@1/worker@1"); + } + + @Test + public void runNode_itemFails_cancelsRunningSiblings() { + AtomicBoolean siblingCancelled = new AtomicBoolean(); + Node worker = + node( + item -> + item.equals("slow") + ? Flowable.timer(10, SECONDS).doOnCancel(() -> siblingCancelled.set(true)) + : Flowable.error(new IllegalStateException("boom"))); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of("slow", "bad")); + + assertThat(run.failure().orElseThrow()).hasMessageThat().isEqualTo("boom"); + assertThat(siblingCancelled.get()).isTrue(); + } + + @Test + public void runNode_itemThrowsAnError_propagatesItAndCancelsRunningSiblings() { + AtomicBoolean siblingCancelled = new AtomicBoolean(); + Node worker = + node( + item -> + item.equals("slow") + ? Flowable.timer(10, SECONDS).doOnCancel(() -> siblingCancelled.set(true)) + : Flowable.error(new AssertionError("fatal"))); + Node node = ParallelWorkers.parallelNode(worker); + + AssertionError thrown = + assertThrows(AssertionError.class, () -> run(node, ImmutableList.of("slow", "fatal"))); + + assertThat(thrown).hasMessageThat().isEqualTo("fatal"); + assertThat(siblingCancelled.get()).isTrue(); + } + + @Test + public void runNode_itemInterrupts_waitsWithoutOutput() { + Run run = run(ParallelWorkers.parallelNode(interruptsOn(2)), ImmutableList.of(1, 2)); + + assertThat(run.interruptIds()).containsExactly("ask-2"); + assertThat(run.output()).isNull(); + assertThat(run.failure()).isEmpty(); + } + + @Test + public void runNode_itemInterrupts_letsRunningItemsFinish() { + CountDownLatch release = new CountDownLatch(1); + Node worker = + node( + item -> + item.equals("slow") + ? Flowable.fromCallable( + () -> { + awaitRelease(release); + return "done"; + }) + .subscribeOn(Schedulers.io()) + : Flowable.just(interrupt("ask-" + item))); + List events = new CopyOnWriteArrayList<>(); + TestObserver observer = + start(ParallelWorkers.parallelNode(worker), ImmutableList.of("slow", "asks"), events) + .test(); + + release.countDown(); + + NodeExecutionState state = observer.awaitDone(10, SECONDS).values().get(0).requireNodeState(); + assertThat(state.interruptIds()).containsExactly("ask-asks"); + assertThat(state.output()).isEmpty(); + assertThat(outputEventAt(events, "worker@1/worker@1").output()).hasValue("done"); + } + + @Test + public void runNode_itemInterruptsUnderCap_neverStartsQueuedItems() { + Run run = + run( + ParallelWorkers.parallelNode(interruptsOn(1), /* maxParallelWorkers= */ 1), + ImmutableList.of(1, 2, 3)); + + assertThat(run.interruptIds()).containsExactly("ask-1"); + assertThat(ImmutableSet.copyOf(pathsOf(run.eventsUnder("worker@1/")))) + .containsExactly("worker@1/worker@1"); + } + + @Test + public void runNode_siblingFailsAfterAnInterrupt_failsTheNode() { + Node worker = + node( + item -> + item.equals(1) + ? Flowable.just(interrupt("ask-1")) + : Flowable.error(new IllegalStateException("late failure"))); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2)); + + assertThat(run.failure().orElseThrow()).hasMessageThat().isEqualTo("late failure"); + } + + @Test + public void runNode_severalItemsInterrupt_mergesAllInterruptIds() { + Node worker = node(item -> Flowable.just(interrupt("ask-" + item))); + + Run run = run(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2)); + + assertThat(run.interruptIds()).containsExactly("ask-1", "ask-2"); + } + + @Test + public void resume_answeredItem_replaysTheFinishedItemsWithoutRunningThemAgain() { + Map runsByItem = new ConcurrentHashMap<>(); + Node worker = + node( + item -> { + runsByItem.merge(item, 1, Integer::sum); + return item.equals(2) ? Flowable.just(interrupt("ask-2")) : Flowable.just(item); + }); + Node items = emitter("items", ImmutableList.of(1, 2, 3)); + Node fanOut = ParallelWorkers.parallelNode(worker); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, items), + new Edge(items, fanOut), + new Edge(fanOut, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask-2", ImmutableMap.of("result", 20))); + + assertThat(runsByItem).containsExactly(1, 1, 2, 1, 3, 1); + assertThat(replyInput(events)).isEqualTo(ImmutableList.of(1, 20, 3)); + } + + @Test + public void resume_underCap_startsTheQueuedItemOnceTheAskingItemIsAnswered() { + Map runsByItem = new ConcurrentHashMap<>(); + Node worker = + node( + item -> { + runsByItem.merge(item, 1, Integer::sum); + return item.equals(1) ? Flowable.just(item) : Flowable.just(interrupt("ask-" + item)); + }); + Node items = emitter("items", ImmutableList.of(1, 2, 3)); + Node fanOut = ParallelWorkers.parallelNode(worker, /* maxParallelWorkers= */ 2); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, items), + new Edge(items, fanOut), + new Edge(fanOut, captureInput("reply")))); + session.advance(go()); + ImmutableMap runsAfterFirstTurn = ImmutableMap.copyOf(runsByItem); + session.advance(answer("ask-2", ImmutableMap.of("result", 20))); + + ImmutableList events = session.turn(answer("ask-3", ImmutableMap.of("result", 30))); + + assertThat(runsAfterFirstTurn).containsExactly(1, 1, 2, 1); + assertThat(runsByItem).containsExactly(1, 1, 2, 1, 3, 1); + assertThat(replyInput(events)).isEqualTo(ImmutableList.of(1, 20, 30)); + } + + @Test + public void runNode_disposed_cancelsRunningItems() throws Exception { + CountDownLatch started = new CountDownLatch(2); + CountDownLatch cancelled = new CountDownLatch(2); + Node worker = + node( + item -> + Flowable.timer(10, SECONDS) + .doOnSubscribe(unused -> started.countDown()) + .doOnCancel(cancelled::countDown)); + TestObserver observer = + start(ParallelWorkers.parallelNode(worker), ImmutableList.of(1, 2)).test(); + assertThat(started.await(10, SECONDS)).isTrue(); + + observer.dispose(); + + observer.assertNoErrors(); + assertThat(cancelled.await(10, SECONDS)).isTrue(); + } + + @Test + public void runNode_textContentWithJsonList_fansOutOverItsElements() { + Run run = + run(ParallelWorkers.parallelNode(doubler()), Content.fromParts(Part.fromText("[1, 2, 3]"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of(2, 4, 6)); + } + + @Test + public void runNode_textContentWithJsonScalar_runsOnTheParsedValue() { + Run run = run(ParallelWorkers.parallelNode(doubler()), Content.fromParts(Part.fromText("5"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of(10)); + } + + @Test + public void runNode_textThatIsNotJson_runsAsOneRawTextItem() { + Run run = run(ParallelWorkers.parallelNode(echo()), Content.fromParts(Part.fromText("hello"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of("hello")); + } + + @Test + public void runNode_contentWithNonTextPart_runsAsOneItem() { + Content message = + Content.fromParts(Part.fromText("[1, 2]"), Part.fromFunctionCall("f", ImmutableMap.of())); + + Run run = run(ParallelWorkers.parallelNode(echo()), message); + + assertThat(run.output()).isEqualTo(ImmutableList.of(message)); + } + + @Test + public void runNode_contentWithAThoughtBeforeAJsonList_fansOutOverTheList() { + Content message = + Content.fromParts( + Part.builder().text("Listing the numbers.").thought(true).build(), + Part.fromText("[1, 2, 3]")); + + Run run = run(ParallelWorkers.parallelNode(doubler()), message); + + assertThat(run.output()).isEqualTo(ImmutableList.of(2, 4, 6)); + } + + @Test + public void runNode_jsonObjectTextForAStringItem_staysRawText() { + Node worker = echoTaking(Schema.builder().type("STRING").build()); + + Run run = + run(ParallelWorkers.parallelNode(worker), Content.fromParts(Part.fromText("{\"a\":1}"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of("{\"a\":1}")); + } + + @Test + public void runNode_jsonObjectTextForAnObjectItem_getsTheParsedMap() { + Node worker = echoTaking(Schema.builder().type("OBJECT").build()); + + Run run = + run(ParallelWorkers.parallelNode(worker), Content.fromParts(Part.fromText("{\"a\":1}"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of(ImmutableMap.of("a", 1))); + } + + @Test + public void runNode_jsonObjectTextForAnItemTakingAStringThroughAnyOf_staysRawText() { + Node worker = + echoTaking( + Schema.builder() + .anyOf( + ImmutableList.of( + Schema.builder().type("INTEGER").build(), + Schema.builder().type("STRING").build())) + .build()); + + Run run = + run(ParallelWorkers.parallelNode(worker), Content.fromParts(Part.fromText("{\"a\":1}"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of("{\"a\":1}")); + } + + @Test + public void runNode_jsonListTextForAStringItem_fansOutOverTheList() { + Node worker = echoTaking(Schema.builder().type("STRING").build()); + + Run run = + run( + ParallelWorkers.parallelNode(worker), + Content.fromParts(Part.fromText("[\"a\", \"b\"]"))); + + assertThat(run.output()).isEqualTo(ImmutableList.of("a", "b")); + } + + /** The outcome of running a node as the root of a workflow run. */ + private record Run(Context context, ImmutableList events) { + @Nullable Object output() { + return context.requireNodeState().output().orElse(null); + } + + Optional failure() { + return context.requireNodeState().failure(); + } + + ImmutableSet interruptIds() { + return context.requireNodeState().interruptIds(); + } + + ImmutableList eventsUnder(String pathPrefix) { + return events.stream() + .filter(event -> pathOf(event).startsWith(pathPrefix)) + .collect(toImmutableList()); + } + } + + private static Run run(Node node, @Nullable Object nodeInput) { + List events = new CopyOnWriteArrayList<>(); + Context context = start(node, nodeInput, events).blockingGet(); + return new Run(context, ImmutableList.copyOf(events)); + } + + /** Returns the run of {@code node} as the root of a workflow run, ignoring its events. */ + private static Single start(Node node, @Nullable Object nodeInput) { + return start(node, nodeInput, new CopyOnWriteArrayList<>()); + } + + /** + * Returns the run of {@code node} as the root of a workflow run, adding its events to {@code + * events}. + */ + private static Single start(Node node, @Nullable Object nodeInput, List events) { + return NodeRunner.builder() + .node(node) + .invocationContext(TestUtils.createInvocationContext(TestUtils.createRootAgent())) + .eventSink((event, nodeFailure) -> Completable.fromAction(() -> events.add(event))) + .build() + .run(nodeInput); + } + + /** Returns a plain node named "worker", not a {@link BaseNode}, that runs {@code body}. */ + private static Node node(Function> body) { + return new TestNode("worker", (context, item) -> Flowable.defer(() -> body.apply(item))); + } + + /** Returns a node named "worker" with {@code config} that runs {@code body}. */ + private static Node node(NodeConfig config, Function> body) { + return new Node() { + @Override + public String name() { + return "worker"; + } + + @Override + public NodeConfig config() { + return config; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.defer(() -> body.apply(nodeInput)); + } + }; + } + + /** Returns a node named "worker" that takes {@code inputSchema} and outputs its input. */ + private static Node echoTaking(Schema inputSchema) { + return new Node() { + @Override + public String name() { + return "worker"; + } + + @Override + public Optional inputSchema() { + return Optional.of(inputSchema); + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.just(nodeInput); + } + }; + } + + private static Node doubler() { + return node(item -> Flowable.just((Integer) item * 2)); + } + + private static Node echo() { + return node(Flowable::just); + } + + private static Node failsOn(Object bad) { + return node( + item -> + item.equals(bad) + ? Flowable.error(new IllegalStateException("bad item " + item)) + : Flowable.just(item)); + } + + private static Node interruptsOn(Object asking) { + return node( + item -> + item.equals(asking) ? Flowable.just(interrupt("ask-" + item)) : Flowable.just(item)); + } + + /** Waits for {@code latch} to be released, failing the item if it never is. */ + private static void awaitRelease(CountDownLatch latch) throws InterruptedException { + if (!latch.await(10, SECONDS)) { + throw new IllegalStateException("The latch was not released."); + } + } +}