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

Rolled back due to a breakage.

PiperOrigin-RevId: 1005597242

Author: Mateusz Krawiec

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: one other than its sub-agents, parent, and peers; one its {@code disallowTransferToParent} or
- * {@code disallowTransferToPeers} setting excludes; or itself, outside a live run.
+ * to: itself, an agent other than its sub-agents, parent, and peers, or one its {@code
+ * disallowTransferToParent} or {@code disallowTransferToPeers} setting excludes.
  */
 public final class TransferRejectedException extends IllegalStateException {
 

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

@@ -583,8 +583,7 @@ 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, which fails with {@link
-   *     TransferRejectedException} if an agent in the run requests a transfer it may not make.
+   * @return A {@link Flowable} of {@link Event}s streamed in real-time.
    */
   @Override
   public Flowable<Event> runLive(InvocationContext invocationContext) {
@@ -704,7 +703,28 @@ public void onError(Throwable e) {
                                 llmResponse,
                                 callLlmContext);
                           })
-                      .flatMap(event -> followLiveTransfer(event, invocationContext, spanContext))
+                      .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;
+                          })
                       .doOnNext(
                           event -> {
                             ImmutableList<FunctionResponse> functionResponses =
@@ -732,32 +752,6 @@ 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.
    *
@@ -837,16 +831,6 @@ 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,7 +60,6 @@
 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;
@@ -243,138 +242,6 @@ 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"));
@@ -1086,22 +953,4 @@ 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();
-  }
 }