feat: share one run per Context.runNode run ID

PiperOrigin-RevId: 1005539735

Author: damianmomotgoogle

core/src/main/java/com/google/adk/workflow/NodeExecutionState.java

@@ -26,6 +26,7 @@
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import io.reactivex.rxjava3.subjects.SingleSubject;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
@@ -61,6 +62,7 @@ public final class NodeExecutionState {
   private final ImmutableList<String> outputFor;
   private final Map<String, @Nullable Object> resumeInputs;
   private final ConcurrentMap<String, Integer> childRunCounts = new ConcurrentHashMap<>();
+  private final ConcurrentMap<String, SingleSubject<Context>> childRuns = new ConcurrentHashMap<>();
   // The paths of the node runs started but not finished in this turn, shared under one root.
   private final Set<String> unfinishedNodePaths;
   private final Set<String> interruptIds = new CopyOnWriteArraySet<>();
@@ -224,6 +226,14 @@ String nextChildRunId(String name) {
     return String.valueOf(childRunCounts.merge(name, 1, Integer::sum));
   }
 
+  /**
+   * Returns the runs of the children this node runs itself, keyed by {@code name@runId}, which
+   * calls with the same run ID share.
+   */
+  ConcurrentMap<String, SingleSubject<Context>> childRuns() {
+    return childRuns;
+  }
+
   /** Records that this run started; it counts as canceled mid-run until it finishes. */
   void markStarted() {
     unfinishedNodePaths.add(nodePath);
@@ -255,6 +265,11 @@ synchronized ImmutableMap<String, ResumeScan.RecoveredNode> recoveredChildren(
     return recoveredChildren;
   }
 
+  /** Returns the parent's state if this run's output is also the parent's, or null. */
+  @Nullable NodeExecutionState outputParent() {
+    return outputParent;
+  }
+
   /** Returns whether an event has carried the node's own or delegated output. */
   boolean outputEmitted() {
     OutputStatus status = outputState.get().status();

core/src/main/java/com/google/adk/workflow/NodeRunner.java

@@ -40,10 +40,12 @@
 import io.reactivex.rxjava3.processors.FlowableProcessor;
 import io.reactivex.rxjava3.processors.UnicastProcessor;
 import io.reactivex.rxjava3.subjects.CompletableSubject;
+import io.reactivex.rxjava3.subjects.SingleSubject;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ConcurrentMap;
 import java.util.function.Function;
 import org.jspecify.annotations.Nullable;
 
@@ -218,8 +220,7 @@ private static Maybe<Object> startChild(
       parentState.claimOutputDelegation();
     }
     String runId = config.runId().orElseGet(() -> parentState.nextChildRunId(node.name()));
-    // Deferred so that a throw from scan or intercept releases the delegation, as a failure does.
-    return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config))
+    return sharedChildRun(parent, node, nodeInput, runId, config)
         .doOnError(
             error -> {
               if (delegatesOutput) {
@@ -229,6 +230,50 @@ private static Maybe<Object> startChild(
         .flatMapMaybe(child -> childResult(parent, child, delegatesOutput, config.raiseOnWait()));
   }
 
+  /**
+   * Returns the run of {@code node} under {@code runId} shared by the caller's calls with that run
+   * ID, starting it if needed. A run that fails or waits for an output it did not produce is not
+   * shared, so a later call runs the child again. When the call that started a run is canceled, a
+   * call still waiting on it runs the child itself.
+   */
+  private static Single<Context> sharedChildRun(
+      Context parent, Node node, @Nullable Object nodeInput, String runId, RunNodeConfig config) {
+    return Single.defer(
+        () -> {
+          ConcurrentMap<String, SingleSubject<Context>> childRuns =
+              parent.requireNodeState().childRuns();
+          String childKey = segment(node.name(), runId);
+          SingleSubject<Context> run = SingleSubject.create();
+          SingleSubject<Context> shared = childRuns.putIfAbsent(childKey, run);
+          if (shared != null) {
+            return shared.onErrorResumeNext(
+                error ->
+                    error instanceof ChildRunCanceledException
+                        ? sharedChildRun(parent, node, nodeInput, runId, config)
+                        : Single.error(error));
+          }
+          // Defer so a throw from scan or intercept releases the claimed run instead of leaking it.
+          return Single.defer(() -> runOrReplayChild(parent, node, nodeInput, runId, config))
+              .doOnSuccess(
+                  child -> {
+                    if (!keepsItsRun(child)) {
+                      childRuns.remove(childKey, run);
+                    }
+                    run.onSuccess(child);
+                  })
+              .doOnError(
+                  error -> {
+                    childRuns.remove(childKey, run);
+                    run.onError(error);
+                  })
+              .doOnDispose(
+                  () -> {
+                    childRuns.remove(childKey, run);
+                    run.onError(new ChildRunCanceledException());
+                  });
+        });
+  }
+
   /**
    * Runs {@code node} as a child of {@code parent}, or replays the result it recorded in an earlier
    * turn under the same run ID.
@@ -258,6 +303,22 @@ private static Single<Context> runOrReplayChild(
         .run(nodeInput);
   }
 
+  /**
+   * Returns whether a later call with the same run ID shares this finished run of a child: a failed
+   * child runs again, as does one that waits for an output and left no outcome.
+   */
+  private static boolean keepsItsRun(Context child) {
+    NodeExecutionState state = child.requireNodeState();
+    if (state.failure().isPresent()) {
+      return false;
+    }
+    return !state.interruptIds().isEmpty()
+        || state.output().isPresent()
+        || state.selectedRoutes().isPresent()
+        || child.eventActions().transferToAgent().isPresent()
+        || !state.node().waitForOutput();
+  }
+
   /**
    * Returns the child's output, or fails with the child's failure or with an interrupt, which the
    * caller takes on as its own.
@@ -292,9 +353,11 @@ private static Maybe<Object> childResult(
         && (childNode instanceof Workflow || childNode.waitForOutput())) {
       return Maybe.error(new NodeInterruptedException());
     }
-    // No event of this run carries a delegate's replayed output, so the caller emits it.
+    // No event of this run carries a delegate's replayed or shared output, so the caller emits it.
     Optional<Object> output = childState.output();
-    if (delegatesOutput && output.isPresent() && !childState.outputEmitted()) {
+    boolean emittedForCaller =
+        childState.outputEmitted() && childState.outputParent() == parentState;
+    if (delegatesOutput && output.isPresent() && !emittedForCaller) {
       parentState.adoptReplayedOutput(output.get());
     }
     return Maybe.fromOptional(output);
@@ -532,6 +595,13 @@ private Event.Builder newEvent() {
     }
   }
 
+  /** Ends the waiting calls of a shared child run whose first caller was canceled. */
+  private static final class ChildRunCanceledException extends RuntimeException {
+    ChildRunCanceledException() {
+      super("The call that ran this child was canceled.");
+    }
+  }
+
   /**
    * An event queued for the runner, paired with the signal that completes once the caller receives
    * the event and, for an error event, with the failure it reports.

core/src/main/java/com/google/adk/workflow/RunNodeConfig.java

@@ -31,9 +31,10 @@ public abstract class RunNodeConfig {
 
   /**
    * Returns the ID of the child's run, which names the run in node paths and branches. A child that
-   * recorded an output under the same ID in an earlier turn returns it without running again. When
-   * this ID is empty, the child receives the next number among the caller's runs of nodes with the
-   * same name, in call order, so concurrent children need explicit IDs.
+   * already recorded an output under the same ID, in this run of the caller or in an earlier turn,
+   * returns it without running again; one that finished without an output may run again. When this
+   * ID is empty, the child receives the next number among the caller's runs of nodes with the same
+   * name, in call order, so concurrent children need explicit IDs.
    */
   public abstract Optional<String> runId();