fix: restrict transfer_to_agent to sub-agents, parent, and peers

- Rejected transfers throw `IllegalArgumentException` instead of `IllegalStateException`.

PiperOrigin-RevId: 1000022357

Author: kvmilos

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

@@ -679,6 +679,9 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
           if (invocationContext.agentStates().containsKey(name())) {
             Optional<BaseAgent> resumeTarget = findSubAgentToResume(invocationContext);
             if (resumeTarget.isPresent()) {
+              if (endedAfterLatestTransfer(invocationContext, resumeTarget.get())) {
+                return endOfAgentAndRecord(invocationContext);
+              }
               AtomicBoolean resumePaused = new AtomicBoolean(false);
               return resumeTarget
                   .get()
@@ -703,13 +706,21 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
           Flowable<Event> events =
               llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
           AtomicBoolean paused = new AtomicBoolean(false);
+          AtomicBoolean lastEventRejectedTransfer = new AtomicBoolean(false);
           return events
               .doOnNext(
                   event -> {
                     if (invocationContext.shouldPauseInvocation(event)) {
                       paused.set(true);
                     }
+                    lastEventRejectedTransfer.set(isRejectedTransfer(event));
                   })
+              // End this agent before a rejected transfer fails, so a resume does not rerun it.
+              .onErrorResumeNext(
+                  error ->
+                      lastEventRejectedTransfer.get()
+                          ? endOfAgentAndRecord(invocationContext).concatWith(Flowable.error(error))
+                          : Flowable.error(error))
               .concatWith(
                   Flowable.defer(
                       () -> {
@@ -739,14 +750,58 @@ private Optional<String> transferTargetNameFrom(Event event) {
     return event.actions().transferToAgent().filter(target -> !target.equals(name()));
   }
 
+  /**
+   * Returns whether {@code event} is a transfer by this agent that {@code BaseLlmFlow} rejects: one
+   * to itself or to an agent outside {@link #transferTargets()}.
+   */
+  private boolean isRejectedTransfer(Event event) {
+    return name().equals(event.author())
+        && event
+            .actions()
+            .transferToAgent()
+            .filter(
+                target ->
+                    target.equals(name())
+                        || transferTargets().stream()
+                            .noneMatch(agent -> agent.name().equals(target)))
+            .isPresent();
+  }
+
+  /**
+   * Returns whether {@code target} ended after this agent's latest transfer, as it does on its own
+   * rejected transfer. An end marker from before that transfer, such as a past loop iteration, does
+   * not count.
+   */
+  private boolean endedAfterLatestTransfer(InvocationContext context, BaseAgent target) {
+    // All branches: under a ParallelAgent, the target writes its marker on a descendant branch.
+    for (Event event :
+        context.events(/* currentInvocation= */ true, /* currentBranch= */ false).reverse()) {
+      if (transferTargetNameFrom(event).isPresent()) {
+        return false;
+      }
+      if (isEndMarker(event, target.name())) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Returns whether {@code event} is {@code author}'s end-of-agent marker, which has no content.
+   */
+  private static boolean isEndMarker(Event event, String author) {
+    return author.equals(event.author())
+        && event.actions().endOfAgent()
+        && event.content().isEmpty();
+  }
+
   /**
    * Resolves a transfer target recorded in history back to its agent. A live transfer fails in
    * {@code BaseLlmFlow} when the target is gone; a resumed one has no such check, and returning
    * empty here would silently re-run this agent instead of the sub-agent it handed off to.
    */
   private BaseAgent resolveTransferTarget(String targetName) {
-    // Declared targets first, as in Python, so a same-named agent on an unrelated branch is not
-    // resumed into; then the same tree-wide search the live transfer path uses.
+    // Check declared targets first, so a same-named agent elsewhere is not resumed instead.
     for (BaseAgent target : transferTargets()) {
       if (target.name().equals(targetName)) {
         return target;
@@ -793,7 +848,12 @@ private Optional<BaseAgent> findSubAgentToResume(InvocationContext context) {
     }
     // A user event answering no call falls through; Python fails the resume instead.
     for (int i = events.size() - 2; i >= 0; i--) {
-      Optional<String> targetName = transferTargetNameFrom(events.get(i));
+      Event event = events.get(i);
+      // Transfers before this agent's own end marker belong to a finished run, such as a past loop.
+      if (isEndMarker(event, name())) {
+        break;
+      }
+      Optional<String> targetName = transferTargetNameFrom(event);
       if (targetName.isPresent()) {
         return Optional.of(resolveTransferTarget(targetName.get()));
       }

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

@@ -426,7 +426,7 @@ private Single<LlmResponse> handleAfterModelCallback(
    *
    * @return A {@link Flowable} of {@link Event} objects from this step.
    * @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
-   * @throws IllegalStateException if a transfer agent is specified but not found.
+   * @throws IllegalArgumentException if a transfer names an agent this agent may not transfer to.
    */
   private Flowable<Event> runOneStep(Context spanContext, InvocationContext context) {
     AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
@@ -801,34 +801,17 @@ private Flowable<Event> followTransfer(
       return self;
     }
     String agentToTransfer = event.actions().transferToAgent().get();
-    Optional<BaseAgent> nextAgent = context.agent().rootAgent().findAgent(agentToTransfer);
-    if (nextAgent.isEmpty()) {
-      logger.error("Agent not found: {}", agentToTransfer);
-      return self.concatWith(
-          Flowable.error(new IllegalStateException("Agent not found: " + agentToTransfer)));
-    }
-    BaseAgent target;
-    // Resumable only: restricting the default flow too would move behavior for existing callers.
-    if (context.isResumable() && context.agent() instanceof LlmAgent llmAgent) {
-      Optional<BaseAgent> declared = declaredTransferTarget(llmAgent, agentToTransfer);
-      if (declared.isEmpty()) {
-        return self.concatWith(
-            Flowable.error(
-                new IllegalStateException(
-                    "Agent "
-                        + llmAgent.name()
-                        + " is not allowed to transfer to the requested agent.")));
-      }
-      target = declared.get();
-    } else {
-      target = nextAgent.get();
+    LlmAgent agent = (LlmAgent) context.agent();
+    Optional<BaseAgent> target = declaredTransferTarget(agent, agentToTransfer);
+    if (target.isEmpty()) {
+      return self.concatWith(Flowable.error(transferRejected(agent, agentToTransfer)));
     }
     // Otherwise the target can build its request from a session missing the transfer event.
     return self.concatWith(
         PersistBarrier.awaitPersisted(context, ImmutableList.of(event))
             .andThen(
                 // Not around the barrier: andThen subscribes the target on the append's thread.
-                Flowable.defer(() -> target.runAsync(context))
+                Flowable.defer(() -> target.get().runAsync(context))
                     .compose(Tracing.withContext(spanContext))));
   }
 
@@ -839,13 +822,31 @@ private Flowable<Event> followTransfer(
    */
   private static Optional<BaseAgent> declaredTransferTarget(LlmAgent from, String name) {
     if (name.equals(from.name())) {
-      return Optional.of(from);
+      return Optional.empty();
     }
     return AgentTransfer.transferTargets(from).stream()
         .filter(candidate -> candidate.name().equals(name))
         .findFirst();
   }
 
+  /**
+   * Returns the error for a transfer {@link #declaredTransferTarget} rejected. The message names
+   * only {@code from}, since {@code name} is model output.
+   */
+  private static IllegalArgumentException transferRejected(LlmAgent from, String name) {
+    if (name.equals(from.name())) {
+      return new IllegalArgumentException("Agent '" + from.name() + "' cannot transfer to itself.");
+    }
+    if (from.rootAgent().findAgent(name).isEmpty()) {
+      return new IllegalArgumentException(
+          "Agent '"
+              + from.name()
+              + "' requested a transfer to an agent that is not in the agent tree.");
+    }
+    return new IllegalArgumentException(
+        "Agent '" + from.name() + "' is not allowed to transfer to the requested agent.");
+  }
+
   /**
    * Runs {@code callEvent}'s function calls and emits what follows from them: the tool-confirmation
    * request, the function response, and any structured final response. Shared by the normal

core/src/test/java/com/google/adk/agents/LlmAgentTest.java

@@ -835,6 +835,112 @@ public void runAsync_resumeAfterTransfer_userEventWithoutFunctionResponse_resume
     assertThat(rootLlm.getRequests()).isEmpty();
   }
 
+  // The transfer belongs to a run of root that already ended, as in a past loop iteration.
+  @Test
+  public void runAsync_resumeAfterThisAgentEnded_ignoresTheEarlierTransfer() {
+    LlmAgent sub =
+        createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub should not run")))
+            .name("sub")
+            .build();
+    TestLlm rootLlm = createTestLlm(createTextLlmResponse("root runs again"));
+    LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build();
+    InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
+    for (Event event :
+        ImmutableList.of(
+            Event.builder()
+                .id("e1")
+                .invocationId("inv")
+                .author("root")
+                .actions(EventActions.builder().endOfAgent(true).build())
+                .build(),
+            Event.builder()
+                .id("u1")
+                .invocationId("inv")
+                .author("user")
+                .content(Content.fromParts(Part.fromText("next")))
+                .build())) {
+      var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
+    }
+    context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);
+    context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true);
+
+    List<Event> events = root.runAsync(context).toList().blockingGet();
+
+    assertThat(simplifyEvents(events)).contains("root: root runs again");
+    assertThat(simplifyEvents(events)).doesNotContain("sub: sub should not run");
+  }
+
+  // The target's end marker is from before root's latest transfer, as in a past loop iteration.
+  @Test
+  public void runAsync_resumeTransferAfterTargetEndedEarlier_runsTheTarget() {
+    TestLlm subLlm = createTestLlm(createTextLlmResponse("sub runs"));
+    LlmAgent sub = createTestAgentBuilder(subLlm).name("sub").build();
+    LlmAgent root =
+        createTestAgentBuilder(createTestLlm(createTextLlmResponse("root should not run")))
+            .name("root")
+            .subAgents(sub)
+            .build();
+    InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
+    for (Event event :
+        ImmutableList.of(
+            Event.builder()
+                .id("e1")
+                .invocationId("inv")
+                .author("sub")
+                .actions(EventActions.builder().endOfAgent(true).build())
+                .build(),
+            Event.builder()
+                .id("e2")
+                .invocationId("inv")
+                .author("root")
+                .actions(EventActions.builder().endOfAgent(true).build())
+                .build(),
+            transferEvent("root", "sub").toBuilder().id("t2").build())) {
+      var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
+    }
+    context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);
+    context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true);
+
+    List<Event> events = root.runAsync(context).toList().blockingGet();
+
+    assertThat(simplifyEvents(events)).contains("sub: sub runs");
+  }
+
+  // The deprecated setEndInvocation flag on a content event is not an end marker.
+  @Test
+  public void runAsync_resumeAfterContentWithEndFlag_stillResumesTheTransfer() {
+    LlmAgent sub =
+        createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub resumes")))
+            .name("sub")
+            .build();
+    TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run"));
+    LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build();
+    InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
+    for (Event event :
+        ImmutableList.of(
+            Event.builder()
+                .id("e1")
+                .invocationId("inv")
+                .author("root")
+                .content(Content.fromParts(Part.fromText("ending")))
+                .actions(EventActions.builder().endOfAgent(true).build())
+                .build(),
+            Event.builder()
+                .id("u1")
+                .invocationId("inv")
+                .author("user")
+                .content(Content.fromParts(Part.fromText("next")))
+                .build())) {
+      var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
+    }
+    context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);
+
+    List<Event> events = root.runAsync(context).toList().blockingGet();
+
+    assertThat(simplifyEvents(events)).contains("sub: sub resumes");
+    assertThat(rootLlm.getRequests()).isEmpty();
+  }
+
   // Mirrors AgentTransferTest.transferTargets_ordersSubAgentsThenParentThenPeers: the flow keeps a
   // second copy of this rule that the agents package cannot reach, so both copies are pinned.
   @Test