fix(workflow): wait for a canceled FunctionNode body before a retry

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

@@ -32,6 +32,8 @@
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.MaybeEmitter;
+import io.reactivex.rxjava3.schedulers.Schedulers;
+import io.reactivex.rxjava3.subjects.CompletableSubject;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Optional;
@@ -166,7 +168,22 @@ boolean edgeInputSchemaInferred() {
   @Override
   public Flowable<?> runNode(Context context, @Nullable Object nodeInput) {
     // Maybe.create rather than defer, so that canceling the node reaches the body while it runs.
-    return Maybe.<Object>create(emitter -> runBody(context, nodeInput, emitter))
+    return Maybe.<Object>create(
+            emitter -> {
+              // A timeout fails the node only once the body returns, so a retry cannot overlap it.
+              CompletableSubject returned = CompletableSubject.create();
+              NodeExecutionState state = context.requireNodeState();
+              // A body run under another node's context must not replace that node's signal.
+              if (state.node() == this) {
+                // Not on the body's thread, which may be interrupted, nor the Runner's busy pool.
+                state.markBodyStarted(returned.observeOn(Schedulers.computation()));
+              }
+              try {
+                runBody(context, nodeInput, emitter);
+              } finally {
+                returned.onComplete();
+              }
+            })
         .subscribeOn(context.invocationContext().scheduler())
         .toFlowable();
   }

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

@@ -24,8 +24,9 @@
  * The body of a {@link FunctionNode}, which takes the node run's context and its single input; the
  * input may be shared with other nodes, so the body must not mutate it. It runs on the {@code
  * Runner}'s scheduler, so it may block, and a small fixed-size scheduler can deadlock bodies that
- * wait for one another. Canceling the node interrupts the body's thread and drops its result; a
- * body that ignores the interrupt keeps its thread until it returns, and a retry can overlap it.
+ * wait for one another. Canceling the node interrupts the body's thread and drops its result; only
+ * the node's own timeout waits for the body to return, so a body that ignores the interrupt holds
+ * up the retry and the run.
  */
 @Experimental
 @FunctionalInterface

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.core.Completable;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.LinkedHashMap;
@@ -74,6 +75,8 @@ public final class NodeExecutionState {
   private final AtomicReference<OutputState> outputState = new AtomicReference<>(NO_OUTPUT);
   // Set once the node drops its own output because a useAsOutput child holds the output.
   private final AtomicBoolean droppedOwnOutput = new AtomicBoolean();
+  private final AtomicReference<Completable> bodyReturned =
+      new AtomicReference<>(Completable.complete());
   private String eventAuthor;
   private @Nullable ImmutableList<Route> selectedRoutes;
   private boolean routesEmitted;
@@ -298,6 +301,16 @@ void markFinished() {
     }
   }
 
+  /** Records that a body of this run started; {@code returned} completes once it returns. */
+  void markBodyStarted(Completable returned) {
+    bodyReturned.set(returned);
+  }
+
+  /** Completes once the body this run started has returned, or at once if none started. */
+  Completable bodyReturned() {
+    return Completable.defer(bodyReturned::get);
+  }
+
   /**
    * Returns whether this run's child {@code childKey}, a {@code name@runId} segment, started in
    * this turn and has not finished. When read before the child runs again, this means a failing

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

@@ -529,16 +529,19 @@ private Completable withTimeout(Completable run) {
         return run;
       }
       Duration limit = timeout.get();
-      // Fails on the timer's thread: the timed-out attempt may still hold the scheduler's threads.
-      Completable timedOut = Completable.error(() -> new NodeTimeoutException(node.name(), limit));
+      // The timeout only interrupts a body, so the node fails once the body has returned.
+      Completable timedOut =
+          nodeState
+              .bodyReturned()
+              .andThen(Completable.error(() -> new NodeTimeoutException(node.name(), limit)));
       io.opentelemetry.context.Context traceContext = io.opentelemetry.context.Context.current();
-      // Off the caller's thread, so a body that blocks its thread cannot hold off the timeout.
+      // Off the caller's thread, so a body that blocks does not block the caller.
       return Completable.timer(0, NANOSECONDS, invocationContext.scheduler())
           // Unlike subscribeOn, canceling this hop does not interrupt the handling of node events.
           .andThen(run.compose(Tracing.withContext(traceContext)))
           // convert saturates, so a timeout too long to count in nanoseconds waits about 292 years.
           .timeout(NANOSECONDS.convert(limit), NANOSECONDS, timedOut)
-          // The failure arrives on the timer's thread, so it needs the caller's trace context too.
+          // The failure arrives on another thread, so it needs the caller's trace context too.
           .compose(Tracing.withContext(traceContext));
     }
 

core/src/test/java/com/google/adk/workflow/FunctionNodeTest.java

@@ -28,6 +28,7 @@
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.junit.Assert.assertThrows;
 
@@ -55,10 +56,12 @@
 import com.fasterxml.jackson.databind.ser.std.ToStringSerializer;
 import com.fasterxml.jackson.databind.util.StdConverter;
 import com.google.adk.agents.Context;
+import com.google.adk.agents.InvocationContext;
 import com.google.adk.apps.App;
 import com.google.adk.events.Event;
 import com.google.adk.events.EventActions;
 import com.google.adk.events.ToolConfirmation;
+import com.google.adk.plugins.BasePlugin;
 import com.google.adk.runner.Runner;
 import com.google.adk.sessions.Session;
 import com.google.adk.testing.TestNode;
@@ -76,6 +79,7 @@
 import com.google.genai.types.VideoMetadata;
 import io.reactivex.rxjava3.core.Completable;
 import io.reactivex.rxjava3.core.Flowable;
+import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Scheduler;
 import io.reactivex.rxjava3.disposables.Disposable;
 import io.reactivex.rxjava3.plugins.RxJavaPlugins;
@@ -1375,6 +1379,265 @@ public void nodeTimeout_sleepingBody_isInterruptedAndFailsWithNodeTimeoutExcepti
     assertThat(result.errorMessage("NodeTimeoutException")).contains("'sleeper'");
   }
 
+  @Test
+  public void nodeTimeout_bodyIgnoringTheInterrupt_runFailsOnlyAfterItReturns() throws Exception {
+    // The body waits for the run to end, at most 2 s, so a run that ends first records it first.
+    List<String> timeline = Collections.synchronizedList(new ArrayList<>());
+    CountDownLatch runEnded = new CountDownLatch(1);
+    CountDownLatch bodyReturned = new CountDownLatch(1);
+    AtomicBoolean interrupted = new AtomicBoolean();
+    FunctionNode<Content, String> stubborn =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("stubborn")
+            .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build())
+            .body((ctx, message) -> {
+              var unused = Uninterruptibles.awaitUninterruptibly(runEnded, 2, SECONDS);
+              interrupted.set(Thread.currentThread().isInterrupted());
+              timeline.add("body returned");
+              bodyReturned.countDown();
+              return "late";
+            })
+            .build();
+    ExecutorService executor = Executors.newCachedThreadPool();
+
+    RunResult result;
+    try {
+      result = run(chain(stubborn), Schedulers.from(executor));
+      timeline.add("run ended");
+      runEnded.countDown();
+      assertThat(bodyReturned.await(10, SECONDS)).isTrue();
+    } finally {
+      executor.shutdownNow();
+    }
+
+    assertThat(timeline).containsExactly("body returned", "run ended").inOrder();
+    assertThat(interrupted.get()).isTrue();
+    assertThat(result.errorMessage("NodeTimeoutException")).contains("'stubborn'");
+  }
+
+  @Test
+  public void nodeTimeoutWithRetry_bodyIgnoringTheInterrupt_retryStartsOnlyAfterItReturns()
+      throws Exception {
+    // Attempt 1 waits for attempt 2, at most 2 s, so an overlapping retry records its start first.
+    AtomicInteger attempts = new AtomicInteger();
+    List<String> timeline = Collections.synchronizedList(new ArrayList<>());
+    CountDownLatch retryStarted = new CountDownLatch(1);
+    CountDownLatch firstReturned = new CountDownLatch(1);
+    AtomicBoolean firstInterrupted = new AtomicBoolean();
+    FunctionNode<Content, String> stubborn =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("stubborn")
+            .config(
+                NodeConfig.builder()
+                    .timeout(Duration.ofMillis(200))
+                    .retryConfig(
+                        RetryConfig.builder()
+                            .maxAttempts(2)
+                            .initialDelay(Duration.ZERO)
+                            .jitter(0)
+                            .build())
+                    .build())
+            .body((ctx, message) -> {
+              int attempt = attempts.incrementAndGet();
+              timeline.add("start " + attempt);
+              if (attempt == 1) {
+                var unused = Uninterruptibles.awaitUninterruptibly(retryStarted, 2, SECONDS);
+                firstInterrupted.set(Thread.currentThread().isInterrupted());
+                timeline.add("end 1");
+                firstReturned.countDown();
+                return "late";
+              }
+              retryStarted.countDown();
+              return "retried";
+            })
+            .build();
+    ExecutorService executor = Executors.newCachedThreadPool();
+
+    RunResult result;
+    try {
+      result = run(chain(stubborn), Schedulers.from(executor));
+      assertThat(firstReturned.await(10, SECONDS)).isTrue();
+    } finally {
+      executor.shutdownNow();
+    }
+
+    assertThat(timeline).containsExactly("start 1", "end 1", "start 2").inOrder();
+    assertThat(firstInterrupted.get()).isTrue();
+    result.assertCompleted();
+    assertThat(result.outputAt("wf@1/stubborn@1")).isEqualTo("retried");
+  }
+
+  @Test
+  public void nodeTimeoutWithRetry_bodyRunsAnotherNodeUnderItsContext_retryWaitsForTheOuterBody()
+      throws Exception {
+    // "inner" runs under "outer"'s context and returns at once; "outer" then outlives its timeout.
+    AtomicInteger attempts = new AtomicInteger();
+    List<String> timeline = Collections.synchronizedList(new ArrayList<>());
+    CountDownLatch retryStarted = new CountDownLatch(1);
+    CountDownLatch firstReturned = new CountDownLatch(1);
+    FunctionNode<Object, String> inner =
+        FunctionNodes.of("inner", Object.class, String.class, (ctx, in) -> "inner out");
+    FunctionNode<Content, String> outer =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("outer")
+            .config(
+                NodeConfig.builder()
+                    .timeout(Duration.ofMillis(200))
+                    .retryConfig(
+                        RetryConfig.builder()
+                            .maxAttempts(2)
+                            .initialDelay(Duration.ZERO)
+                            .jitter(0)
+                            .build())
+                    .build())
+            .body((ctx, message) -> {
+              int attempt = attempts.incrementAndGet();
+              timeline.add("start " + attempt);
+              if (attempt == 1) {
+                inner.runNode(ctx, null).blockingSubscribe();
+                timeline.add("inner returned");
+                var unused = Uninterruptibles.awaitUninterruptibly(retryStarted, 2, SECONDS);
+                timeline.add("end 1");
+                firstReturned.countDown();
+                return "late";
+              }
+              retryStarted.countDown();
+              return "retried";
+            })
+            .build();
+    ExecutorService executor = Executors.newCachedThreadPool();
+
+    RunResult result;
+    try {
+      result = run(chain(outer), Schedulers.from(executor));
+      assertThat(firstReturned.await(10, SECONDS)).isTrue();
+    } finally {
+      executor.shutdownNow();
+    }
+
+    assertThat(timeline).containsExactly("start 1", "inner returned", "end 1", "start 2").inOrder();
+    result.assertCompleted();
+    assertThat(result.outputAt("wf@1/outer@1")).isEqualTo("retried");
+  }
+
+  @Test
+  public void nodeTimeout_oneThreadRunnerWithABusySibling_failsTheRunAndCancelsTheSibling()
+      throws Exception {
+    // Once "sleeper" times out, a task it queues and "busy" hold the one thread until released.
+    // Both nodes have a timeout, so both start through the same hops and "sleeper" starts first.
+    List<String> timeline = Collections.synchronizedList(new ArrayList<>());
+    CountDownLatch release = new CountDownLatch(1);
+    ExecutorService executor = Executors.newSingleThreadExecutor();
+    FunctionNode<Content, String> sleeper =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("sleeper")
+            .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build())
+            .body((ctx, message) -> {
+              timeline.add("sleeper started");
+              try {
+                Thread.sleep(60_000);
+              } catch (InterruptedException e) {
+                executor.execute(() -> Uninterruptibles.awaitUninterruptibly(release));
+              }
+              return "slept";
+            })
+            .build();
+    FunctionNode<Content, String> busy =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("busy")
+            .config(NodeConfig.builder().timeout(Duration.ofHours(1)).build())
+            .body((ctx, message) -> {
+              timeline.add("busy started");
+              release.await();
+              timeline.add("busy finished");
+              return "done";
+            })
+            .build();
+    Workflow workflow =
+        Workflow.builder()
+            .name("wf")
+            .edges(new Edge(Start.INSTANCE, sleeper), new Edge(Start.INSTANCE, busy))
+            .build();
+    Runner runner = runner(workflow, Schedulers.from(executor));
+    Session session = runner.sessionService().createSession(APP, USER).blockingGet();
+    TestSubscriber<Event> subscriber = new TestSubscriber<>();
+
+    boolean runEnded;
+    try {
+      runner
+          .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go")))
+          .subscribe(subscriber);
+      runEnded = subscriber.await(5, SECONDS);
+    } finally {
+      subscriber.cancel();
+      release.countDown();
+      executor.shutdownNow();
+    }
+
+    assertThat(timeline).isNotEmpty();
+    assertThat(timeline.get(0)).isEqualTo("sleeper started");
+    assertWithMessage("whether the run ended within 5 s").that(runEnded).isTrue();
+    assertThat(timeline).doesNotContain("busy finished");
+    subscriber.assertError(NodeTimeoutException.class);
+  }
+
+  @Test
+  public void
+      inlineSchedulerWithInterruptedCaller_nodeTimeout_deliversTheFailureOnAnUninterruptedThread()
+          throws Exception {
+    // The body outlives its timeout on the caller's interrupted thread.
+    AtomicReference<Boolean> errorOnInterruptedThread = new AtomicReference<>();
+    FunctionNode<Content, String> stubborn =
+        FunctionNodes.builder(Content.class, String.class)
+            .name("stubborn")
+            .config(NodeConfig.builder().timeout(Duration.ofMillis(200)).build())
+            .body((ctx, message) -> {
+              Uninterruptibles.sleepUninterruptibly(400, MILLISECONDS);
+              return "late";
+            })
+            .build();
+    BasePlugin observer =
+        new BasePlugin("observer") {
+          @Override
+          public Maybe<Event> onEventCallback(InvocationContext context, Event event) {
+            if (event.errorCode().isPresent()) {
+              errorOnInterruptedThread.set(Thread.currentThread().isInterrupted());
+            }
+            return Maybe.empty();
+          }
+        };
+    Runner runner =
+        Runner.builder()
+            .app(App.builder().name(APP).rootNode(chain(stubborn)).plugins(observer).build())
+            .scheduler(Schedulers.trampoline())
+            .build();
+    Session session = runner.sessionService().createSession(APP, USER).blockingGet();
+    TestSubscriber<Event> subscriber = new TestSubscriber<>();
+    ExecutorService caller = Executors.newSingleThreadExecutor();
+
+    boolean runEnded;
+    try {
+      var unused =
+          caller.submit(() -> {
+            Thread.currentThread().interrupt();
+            runner
+                .runAsync(USER, session.id(), Content.fromParts(Part.fromText("go")))
+                .subscribe(subscriber);
+            return Thread.interrupted();
+          });
+      runEnded = subscriber.await(5, SECONDS);
+    } finally {
+      subscriber.cancel();
+      caller.shutdownNow();
+    }
+
+    assertWithMessage("whether the run ended within 5 s").that(runEnded).isTrue();
+    subscriber.assertError(NodeTimeoutException.class);
+    assertWithMessage("whether the error event reached plugins on an interrupted thread")
+        .that(errorOnInterruptedThread.get())
+        .isFalse();
+  }
+
   @Test
   public void runDisposed_bodyThatThrowsAfterTheInterrupt_reportsNoUndeliverableError()
       throws InterruptedException {