From f4d6de4cc2e4d6765449753e1a3ddf94c579a7db Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Thu, 22 Oct 2026 17:02:43 -0700 Subject: [PATCH] fix: restrict live `transfer_to_agent` to the agent, its sub-agents, parent, and peers PiperOrigin-RevId: 1005160552 --- .../adk/agents/TransferRejectedException.java | 4 +- .../adk/flows/llmflows/BaseLlmFlow.java | 62 ++++--- .../adk/flows/llmflows/AgentTransferTest.java | 151 ++++++++++++++++++ 3 files changed, 192 insertions(+), 25 deletions(-) diff --git a/core/src/main/java/com/google/adk/agents/TransferRejectedException.java b/core/src/main/java/com/google/adk/agents/TransferRejectedException.java index 51adc2dc5..183ab72fb 100644 --- a/core/src/main/java/com/google/adk/agents/TransferRejectedException.java +++ b/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 { diff --git a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java index e53283e54..65b248f61 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java @@ -583,7 +583,8 @@ private Flowable runStepAndCheckContinue( *

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 runLive(InvocationContext invocationContext) { @@ -703,28 +704,7 @@ public void onError(Throwable e) { llmResponse, callLlmContext); }) - .flatMap( - event -> { - Flowable events = Flowable.just(event); - if (event.actions().transferToAgent().isPresent()) { - BaseAgent rootAgent = invocationContext.agent().rootAgent(); - Optional nextAgent = - rootAgent.findAgent(event.actions().transferToAgent().get()); - if (nextAgent.isEmpty()) { - throw new IllegalStateException( - "Agent not found: " + event.actions().transferToAgent().get()); - } - Flowable 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 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 followLiveTransfer( + Event event, InvocationContext context, Context spanContext) { + Flowable self = Flowable.just(event); + if (event.actions().transferToAgent().isEmpty()) { + return self; + } + String agentToTransfer = event.actions().transferToAgent().get(); + LlmAgent agent = (LlmAgent) context.agent(); + Optional 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 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 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. diff --git a/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java b/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java index 702430ff2..3dc255100 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/AgentTransferTest.java +++ b/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 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 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 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 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 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 runRunner(Runner runner, InvocationContext invocationContext .toList() .blockingGet(); } + + private static ImmutableList runLive(LlmAgent root) { + List 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 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(); + } }