fix: restrict live transfer_to_agent to the agent, its sub-agents, parent, and peers

PiperOrigin-RevId: 1005160552

Author: kvmilos

core/src/main/java/com/google/adk/agents/TransferRejectedException.java

@@ -18,8 +18,8 @@
 
 /**
  * Thrown when an {@link LlmAgent} calls {@code transfer_to_agent} with an agent it may not transfer
- * to: itself, an agent other than its sub-agents, parent, and peers, or one its {@code
- * disallowTransferToParent} or {@code disallowTransferToPeers} setting excludes.
+ * to: one other than its sub-agents, parent, and peers; one its {@code disallowTransferToParent} or
+ * {@code disallowTransferToPeers} setting excludes; or itself, outside a live run.
  */
 public final class TransferRejectedException extends IllegalStateException {
 

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -583,7 +583,8 @@ private Flowable<Event> runStepAndCheckContinue(
    * <p>Handles sending history and live requests to the LLM, receiving responses, processing them,
    * and managing agent transfers.
    *
-   * @return A {@link Flowable} of {@link Event}s streamed in real-time.
+   * @return A {@link Flowable} of {@link Event}s streamed in real-time, which fails with {@link
+   *     TransferRejectedException} if an agent in the run requests a transfer it may not make.
    */
   @Override
   public Flowable<Event> runLive(InvocationContext invocationContext) {
@@ -703,28 +704,7 @@ public void onError(Throwable e) {
                                 llmResponse,
                                 callLlmContext);
                           })
-                      .flatMap(
-                          event -> {
-                            Flowable<Event> events = Flowable.just(event);
-                            if (event.actions().transferToAgent().isPresent()) {
-                              BaseAgent rootAgent = invocationContext.agent().rootAgent();
-                              Optional<BaseAgent> nextAgent =
-                                  rootAgent.findAgent(event.actions().transferToAgent().get());
-                              if (nextAgent.isEmpty()) {
-                                throw new IllegalStateException(
-                                    "Agent not found: " + event.actions().transferToAgent().get());
-                              }
-                              Flowable<Event> nextAgentEvents =
-                                  Flowable.defer(
-                                      () -> {
-                                        try (Scope scope = spanContext.makeCurrent()) {
-                                          return nextAgent.get().runLive(invocationContext);
-                                        }
-                                      });
-                              events = Flowable.concat(events, nextAgentEvents);
-                            }
-                            return events;
-                          })
+                      .flatMap(event -> followLiveTransfer(event, invocationContext, spanContext))
                       .doOnNext(
                           event -> {
                             ImmutableList<FunctionResponse> functionResponses =
@@ -752,6 +732,32 @@ public void onError(Throwable e) {
             }));
   }
 
+  /**
+   * Emits {@code event}, followed by the target agent's live run when it carries a transfer. If the
+   * agent may not make that transfer, fails with {@link TransferRejectedException} after {@code
+   * event}, so the caller's {@code doOnNext} still sees the transfer and closes the connection.
+   */
+  private static Flowable<Event> followLiveTransfer(
+      Event event, InvocationContext context, Context spanContext) {
+    Flowable<Event> self = Flowable.just(event);
+    if (event.actions().transferToAgent().isEmpty()) {
+      return self;
+    }
+    String agentToTransfer = event.actions().transferToAgent().get();
+    LlmAgent agent = (LlmAgent) context.agent();
+    Optional<BaseAgent> target = liveTransferTarget(agent, agentToTransfer);
+    if (target.isEmpty()) {
+      return self.concatWith(Flowable.error(transferRejected(agent, agentToTransfer)));
+    }
+    return self.concatWith(
+        Flowable.defer(
+            () -> {
+              try (Scope scope = spanContext.makeCurrent()) {
+                return target.get().runLive(context);
+              }
+            }));
+  }
+
   /**
    * Builds an {@link Event} from LLM response, request, and base event data.
    *
@@ -831,6 +837,16 @@ private static Optional<BaseAgent> declaredTransferTarget(LlmAgent from, String
         .findFirst();
   }
 
+  /**
+   * Returns the agent a live transfer from {@code from} hands over to, or empty when the transfer
+   * is not allowed. Unlike {@link #declaredTransferTarget}, returns {@code from} itself when {@code
+   * name} names it, so a live run reopens it on a new connection; any other name must be a declared
+   * target.
+   */
+  private static Optional<BaseAgent> liveTransferTarget(LlmAgent from, String name) {
+    return name.equals(from.name()) ? Optional.of(from) : declaredTransferTarget(from, name);
+  }
+
   /**
    * Returns the error for a transfer {@link #declaredTransferTarget} rejected. The message names
    * only {@code from}, since {@code name} is model output.

core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java

@@ -60,6 +60,7 @@
 import io.reactivex.rxjava3.core.Flowable;
 import io.reactivex.rxjava3.core.Single;
 import io.reactivex.rxjava3.subscribers.TestSubscriber;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -242,6 +243,138 @@ public void runLive_transferToAgent_closesConnection() throws Exception {
     assertThat(closedConnectionsCount).isEqualTo(1);
   }
 
+  @Test
+  public void runLive_transferToGrandchild_failsAndGrandchildNeverRuns() {
+    TestLlm grandchildLlm = TestLlm.create(createTextLlmResponse("grandchild ran"));
+    LlmAgent root = rootTransferringToGrandchild(grandchildLlm);
+
+    TransferRejectedException error =
+        assertThrows(TransferRejectedException.class, () -> runLive(root));
+
+    assertThat(error)
+        .hasMessageThat()
+        .isEqualTo("Agent 'root' is not allowed to transfer to the requested agent.");
+    assertThat(grandchildLlm.getRequests()).isEmpty();
+  }
+
+  @Test
+  public void runLive_transferToParent_withDisallowTransferToParent_fails() {
+    // A second connection, so a wrongly re-run root ends the run instead of transferring again.
+    TestLlm rootLlm =
+        createTestLlm(
+            Flowable.just(transferTo("sub")),
+            Flowable.just(createTextLlmResponse("root ran again")));
+    LlmAgent root = rootWithSubDisallowingParent(rootLlm);
+
+    TransferRejectedException error =
+        assertThrows(TransferRejectedException.class, () -> runLive(root));
+
+    assertThat(error)
+        .hasMessageThat()
+        .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent.");
+    assertThat(rootLlm.getRequests()).hasSize(1);
+  }
+
+  @Test
+  public void runLive_transferToPeer_withDisallowTransferToPeers_failsAndPeerNeverRuns() {
+    TestLlm peerLlm = TestLlm.create(ImmutableList.of());
+    LlmAgent root = rootWithSubDisallowingPeers(TestLlm.create(transferTo("peer")), peerLlm);
+
+    TransferRejectedException error =
+        assertThrows(TransferRejectedException.class, () -> runLive(root));
+
+    assertThat(error)
+        .hasMessageThat()
+        .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent.");
+    assertThat(error.agentName()).isEqualTo("sub");
+    assertThat(peerLlm.getRequests()).isEmpty();
+  }
+
+  // The model chose the name, so the error names only the calling agent.
+  @Test
+  public void runLive_transferToUnknownAgent_failsAsNotInTreeAfterClosingTheConnection() {
+    TestLlm rootLlm = TestLlm.create(transferTo("missing"));
+    LlmAgent root = root(rootLlm, textAgent("sub"));
+    List<Event> events = new ArrayList<>();
+
+    TransferRejectedException error =
+        assertThrows(TransferRejectedException.class, () -> runLive(root, events));
+
+    assertThat(error)
+        .hasMessageThat()
+        .isEqualTo("Agent 'root' requested a transfer to an agent that is not in the agent tree.");
+    assertThat(simplifyEvents(events))
+        .containsExactly(
+            "root: FunctionCall(name=transfer_to_agent, args={agent_name=missing})",
+            "root: FunctionResponse(name=transfer_to_agent, response={})")
+        .inOrder();
+    assertThat(rootLlm.getLiveRequestHistory().stream().filter(LiveRequest::shouldClose).count())
+        .isEqualTo(1);
+  }
+
+  @Test
+  public void runLive_transferToDuplicateName_runsTheDeclaredTarget() {
+    TestLlm declaredLlm = TestLlm.create(createTextLlmResponse("declared ran"));
+    TestLlm namesakeLlm = TestLlm.create(createTextLlmResponse("namesake ran"));
+
+    ImmutableList<Event> events = runLive(rootWithNamesakeTargets(declaredLlm, namesakeLlm));
+
+    assertThat(simplifyEvents(events)).contains("shared: declared ran");
+    assertThat(declaredLlm.getRequests()).hasSize(1);
+    assertThat(namesakeLlm.getRequests()).isEmpty();
+  }
+
+  // Peers and the parent are declared targets too, not only sub-agents.
+  @Test
+  public void runLive_transferToPeerThenParent_runsBoth() {
+    TestLlm rootLlm =
+        createTestLlm(
+            Flowable.just(transferTo("sub")), Flowable.just(createTextLlmResponse("root again")));
+    TestLlm peerLlm = TestLlm.create(transferTo("root"));
+    LlmAgent sub = createTestAgentBuilder(TestLlm.create(transferTo("peer"))).name("sub").build();
+    LlmAgent peer = createTestAgentBuilder(peerLlm).name("peer").build();
+
+    ImmutableList<Event> events = runLive(root(rootLlm, sub, peer));
+
+    assertThat(simplifyEvents(events)).contains("root: root again");
+    assertThat(peerLlm.getRequests()).hasSize(1);
+    assertThat(rootLlm.getRequests()).hasSize(2);
+  }
+
+  // runAsync rejects a transfer to self; a live run reopens the agent on a new connection.
+  @Test
+  public void runLive_transferToSelf_reopensTheAgent() {
+    TestLlm rootLlm =
+        createTestLlm(
+            Flowable.just(transferTo("root")),
+            Flowable.just(createTextLlmResponse("root reopened")));
+
+    ImmutableList<Event> events = runLive(root(rootLlm, textAgent("sub")));
+
+    assertThat(simplifyEvents(events))
+        .containsExactly(
+            "root: FunctionCall(name=transfer_to_agent, args={agent_name=root})",
+            "root: FunctionResponse(name=transfer_to_agent, response={})",
+            "root: root reopened")
+        .inOrder();
+    assertThat(rootLlm.getRequests()).hasSize(2);
+  }
+
+  @Test
+  public void runLive_transferToSelf_reopensTheCallerNotANamesake() {
+    TestLlm callerLlm =
+        createTestLlm(
+            Flowable.just(transferTo("shared")),
+            Flowable.just(createTextLlmResponse("caller reopened")));
+    TestLlm namesakeLlm = TestLlm.create(createTextLlmResponse("namesake ran"));
+
+    ImmutableList<Event> events = runLive(rootWithNamesakeTargets(callerLlm, namesakeLlm));
+
+    assertThat(simplifyEvents(events)).contains("shared: caller reopened");
+    assertThat(callerLlm.getRequests()).hasSize(2);
+    assertThat(namesakeLlm.getRequests()).isEmpty();
+  }
+
   @Test
   public void testAutoToAuto() {
     Content transferCallContent = Content.fromParts(createTransferCallPart("sub_agent_1"));
@@ -953,4 +1086,22 @@ private List<Event> runRunner(Runner runner, InvocationContext invocationContext
         .toList()
         .blockingGet();
   }
+
+  private static ImmutableList<Event> runLive(LlmAgent root) {
+    List<Event> events = new ArrayList<>();
+    runLive(root, events);
+    return ImmutableList.copyOf(events);
+  }
+
+  /** Runs {@code root} live after one user message, adding each event to {@code events}. */
+  private static void runLive(LlmAgent root, List<Event> events) {
+    Runner runner = new InMemoryRunner(root, "test");
+    LiveRequestQueue liveRequestQueue = new LiveRequestQueue();
+    liveRequestQueue.content(Content.fromParts(Part.fromText("hi")));
+    runner
+        .runLive(newSession(runner), liveRequestQueue, RunConfig.builder().build())
+        .doOnNext(events::add)
+        .ignoreElements()
+        .blockingAwait();
+  }
 }