fix: run approved tool confirmations for agents under a ParallelAgent

PiperOrigin-RevId: 993677550

Author: kvmilos

core/src/main/java/com/google/adk/runner/Runner.java

@@ -27,7 +27,6 @@
 import com.google.adk.agents.InvocationContext;
 import com.google.adk.agents.LiveRequestQueue;
 import com.google.adk.agents.LlmAgent;
-import com.google.adk.agents.ParallelAgent;
 import com.google.adk.agents.Role;
 import com.google.adk.agents.RunConfig;
 import com.google.adk.agents.SequentialAgent;
@@ -416,13 +415,14 @@ private Single<Event> createUserMessageEvent(
       InvocationContext invocationContext,
       boolean saveInputBlobsAsArtifacts,
       @Nullable Map<String, Object> stateDelta) {
+    // As on the resumable path, a function response takes the branch of the call it answers.
     return appendNewMessageToSession(
         session,
         newMessage,
         invocationContext,
         saveInputBlobsAsArtifacts,
         stateDelta,
-        /* branch= */ null);
+        matchingFunctionCallEvent(session, newMessage).flatMap(Event::branch).orElse(null));
   }
 
   private Single<Event> appendNewMessageToSession(
@@ -768,12 +768,14 @@ private Flowable<Event> runAgentForUserEvent(
    */
   private Flowable<Event> runAgentWithUpdatedSession(
       InvocationContext initialContext, Session updatedSession, Event event, BaseAgent rootAgent) {
+    BaseAgent agentToRun = this.findAgentToRun(updatedSession, rootAgent);
     // Create context with updated session for beforeRunCallback
     InvocationContext contextWithUpdatedSession =
         initialContext.toBuilder()
             .session(updatedSession)
-            .agent(this.findAgentToRun(updatedSession, rootAgent))
+            .agent(agentToRun)
             .userContent(event.content().orElseGet(Content::fromParts))
+            .branch(routedAgentParentBranch(updatedSession, agentToRun, rootAgent))
             .build();
 
     // If beforeRunCallback returns content, emit it and skip agent.
@@ -1055,26 +1057,79 @@ private Flowable<Event> runResumedAgent(
   }
 
   /**
-   * Branch to seed a resumed context with so {@code resumeAgent} runs under the same branch it
-   * originally did. Returns the parent branch (the resolved agent's most recent event branch minus
-   * its own trailing name segment, which {@link BaseAgent#runAsync} re-appends), or {@code null}
-   * for the root branch. Non-null only for an agent nested under a {@link ParallelAgent}.
+   * Branch to seed a new invocation with so a routed sub-agent runs where it ran before, or {@code
+   * null} when {@code agentToRun} is the root. A function response restores the branch of the call
+   * it answers, since a Java agent's branch depends on the transfer path that reached it; other
+   * routing restores the agent's latest branch, as Python does.
+   */
+  private static @Nullable String routedAgentParentBranch(
+      Session session, BaseAgent agentToRun, BaseAgent rootAgent) {
+    if (agentToRun.equals(rootAgent)) {
+      return null;
+    }
+    Optional<Event> answeredCall =
+        Functions.findMatchingFunctionCallEvent(session.immutableEvents());
+    if (answeredCall.isPresent()) {
+      String ownSuffix = sequentialBranchSuffix(agentToRun, answeredCall.get().author());
+      if (ownSuffix != null) {
+        return parentBranch(answeredCall.get().branch().orElse(null), ownSuffix);
+      }
+    }
+    return resumeParentBranch(session, /* invocationId= */ null, agentToRun);
+  }
+
+  /**
+   * Branch to seed a context with so {@code resumeAgent} runs under the branch it last ran on: the
+   * parent branch of the newest event by {@code resumeAgent}, or by an agent it reaches through
+   * SequentialAgents, with a non-empty branch. A non-null {@code invocationId} limits the search to
+   * that invocation; {@code null} is returned for the root branch.
    */
   private static @Nullable String resumeParentBranch(
-      Session session, String invocationId, BaseAgent resumeAgent) {
+      Session session, @Nullable String invocationId, BaseAgent resumeAgent) {
     ImmutableList<Event> events = session.immutableEvents();
     for (int i = events.size() - 1; i >= 0; i--) {
       Event event = events.get(i);
-      if (invocationId.equals(event.invocationId())
-          && resumeAgent.name().equals(event.author())
-          && event.branch().isPresent()) {
-        String branch = event.branch().get();
-        String ownSegment = "." + resumeAgent.name();
-        if (branch.endsWith(ownSegment)) {
-          String parent = branch.substring(0, branch.length() - ownSegment.length());
-          return parent.isEmpty() ? null : parent;
+      if ((invocationId == null || invocationId.equals(event.invocationId()))
+          && event.branch().filter(branch -> !branch.isEmpty()).isPresent()) {
+        String ownSuffix = sequentialBranchSuffix(resumeAgent, event.author());
+        if (ownSuffix != null) {
+          return parentBranch(event.branch().get(), ownSuffix);
+        }
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Returns {@code branch} minus the trailing {@code ownSuffix} that {@link BaseAgent#runAsync}
+   * re-appends, or {@code null} for the root branch.
+   */
+  private static @Nullable String parentBranch(@Nullable String branch, String ownSuffix) {
+    if (branch == null || branch.isEmpty() || branch.equals(ownSuffix)) {
+      return null;
+    }
+    if (branch.endsWith("." + ownSuffix)) {
+      String parent = branch.substring(0, branch.length() - ownSuffix.length() - 1);
+      return parent.isEmpty() ? null : parent;
+    }
+    return branch;
+  }
+
+  /**
+   * Returns the dot-joined names from {@code agent} down to the agent named {@code author} when
+   * that agent is {@code agent} or below it through SequentialAgents, else {@code null}. The legacy
+   * flow resumes such an ancestor, which records no events of its own.
+   */
+  private static @Nullable String sequentialBranchSuffix(BaseAgent agent, @Nullable String author) {
+    if (agent.name().equals(author)) {
+      return agent.name();
+    }
+    if (agent instanceof SequentialAgent) {
+      for (BaseAgent subAgent : agent.subAgents()) {
+        String suffix = sequentialBranchSuffix(subAgent, author);
+        if (suffix != null) {
+          return agent.name() + "." + suffix;
         }
-        return branch.equals(resumeAgent.name()) ? null : branch;
       }
     }
     return null;
@@ -1241,11 +1296,13 @@ private InvocationContext newInvocationContextForLive(
         runConfigBuilder.inputAudioTranscription(AudioTranscriptionConfig.builder().build());
       }
     }
+    BaseAgent agentToRun = findAgentToRun(session, this.agent);
     InvocationContext.Builder builder =
-        newInvocationContextBuilder(session, findAgentToRun(session, this.agent))
+        newInvocationContextBuilder(session, agentToRun)
             .runConfig(runConfigBuilder.build())
             .userContent(Content.fromParts())
-            .liveRequestQueue(liveRequestQueue);
+            .liveRequestQueue(liveRequestQueue)
+            .branch(routedAgentParentBranch(session, agentToRun, this.agent));
 
     return builder.build();
   }

core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java

@@ -17,10 +17,14 @@
 package com.google.adk.runner;
 
 import static com.google.adk.testing.ResumabilityTestUtils.answerCall;
+import static com.google.adk.testing.ResumabilityTestUtils.approveConfirmation;
+import static com.google.adk.testing.ResumabilityTestUtils.confirmingEchoFunctionTool;
 import static com.google.adk.testing.ResumabilityTestUtils.newSession;
 import static com.google.adk.testing.ResumabilityTestUtils.pendingFunctionTool;
 import static com.google.adk.testing.ResumabilityTestUtils.runTurn;
+import static com.google.adk.testing.ResumabilityTestUtils.runTurnAskingConfirmation;
 import static com.google.adk.testing.ResumabilityTestUtils.shimRunner;
+import static com.google.adk.testing.ResumabilityTestUtils.textAgent;
 import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse;
 import static com.google.adk.testing.TestUtils.createLlmResponse;
 import static com.google.adk.testing.TestUtils.createTestAgentBuilder;
@@ -475,6 +479,86 @@ public void runAsync_plainTextWithShim_withStateDelta_mergesStateIntoSession() {
     assertThat(finalSession.state()).containsAtLeastEntriesIn(stateDelta);
   }
 
+  // The shim resumes the SequentialAgent, which has no events of its own to restore a branch from.
+  @Test
+  public void
+      runAsync_withToolConfirmation_inSequentialAgentUnderParallelAgent_callsTool_legacyShim() {
+    LlmAgent childAgent =
+        createTestAgentBuilder(
+                createTestLlm(
+                    createFunctionCallLlmResponse(
+                        "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")),
+                    createTextLlmResponse("Response after user confirmed.")))
+            .name("child_agent")
+            .tools(confirmingEchoFunctionTool())
+            .build();
+    SequentialAgent sequentialAgent =
+        SequentialAgent.builder()
+            .name("sequential_agent")
+            .subAgents(ImmutableList.of(childAgent))
+            .build();
+    ParallelAgent rootAgent =
+        ParallelAgent.builder()
+            .name("parallel_agent")
+            .subAgents(
+                ImmutableList.of(sequentialAgent, textAgent("sibling_agent", "Sibling done.")))
+            .build();
+    Runner runner = shimRunner(rootAgent);
+    Session session = newSession(runner);
+    FunctionCall askUserConfirmationFunctionCall =
+        runTurnAskingConfirmation(runner, session, "from user");
+
+    ImmutableList<Event> eventsAfterConfirmation =
+        approveConfirmation(runner, session, askUserConfirmationFunctionCall);
+
+    assertThat(simplifyEvents(eventsAfterConfirmation))
+        .containsExactly(
+            "child_agent: FunctionResponse(name=echoTool, response={message=hello})",
+            "child_agent: Response after user confirmed.")
+        .inOrder();
+    assertThat(eventsAfterConfirmation.stream().map(event -> event.branch().orElse(null)))
+        .containsExactly(
+            "parallel_agent.sequential_agent.child_agent",
+            "parallel_agent.sequential_agent.child_agent");
+  }
+
+  // Answering a long-running call resumes the sequence; its later sub-agent must still run.
+  @Test
+  public void
+      runAsync_withLongRunningCall_inSequentialAgentUnderParallelAgent_runsNextAgent_legacyShim() {
+    LlmAgent childAgent =
+        createTestAgentBuilder(
+                createTestLlm(
+                    createFunctionCallLlmResponse(
+                        "lro_call_id", "pendingTool", ImmutableMap.of("message", "draft")),
+                    createTextLlmResponse("child resumed")))
+            .name("child_agent")
+            .tools(pendingFunctionTool())
+            .build();
+    SequentialAgent sequentialAgent =
+        SequentialAgent.builder()
+            .name("sequential_agent")
+            .subAgents(ImmutableList.of(childAgent, textAgent("next_agent", "next done")))
+            .build();
+    ParallelAgent rootAgent =
+        ParallelAgent.builder()
+            .name("parallel_agent")
+            .subAgents(
+                ImmutableList.of(sequentialAgent, textAgent("sibling_agent", "Sibling done.")))
+            .build();
+    Runner runner = shimRunner(rootAgent);
+    Session session = newSession(runner);
+    var unused = runTurn(runner, session, "from user");
+
+    ImmutableList<Event> eventsAfterResume =
+        answerCall(
+            runner, session, "lro_call_id", "pendingTool", ImmutableMap.of("result", "done"));
+
+    assertThat(simplifyEvents(eventsAfterResume))
+        .containsExactly("child_agent: child resumed", "next_agent: next done")
+        .inOrder();
+  }
+
   // ===== CL1-parity: every CL1 resumable(true) test, re-run under the text-only shim =====
 
   @Test