From 9837e1b15bfb51ef2e20bea13612792f22f0b832 Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Wed, 21 Oct 2026 07:39:01 -0700 Subject: [PATCH] fix: restrict turn-based `transfer_to_agent` to declared targets - Rejected transfers throw the new `TransferRejectedException` (extends `IllegalStateException`). PiperOrigin-RevId: 1003929844 --- .../java/com/google/adk/agents/LlmAgent.java | 93 ++-- .../adk/agents/TransferRejectedException.java | 38 ++ .../adk/flows/llmflows/BaseLlmFlow.java | 55 +-- .../com/google/adk/agents/LlmAgentTest.java | 131 +++++- .../adk/flows/llmflows/AgentTransferTest.java | 413 ++++++++++++++++++ 5 files changed, 674 insertions(+), 56 deletions(-) create mode 100644 core/src/main/java/com/google/adk/agents/TransferRejectedException.java diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index 0bfa3e0f8..d36f0bda0 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -677,8 +677,16 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { () -> { // Resumed after a transfer: continue the sub-agent instead of re-invoking the model. if (invocationContext.agentStates().containsKey(name())) { - Optional resumeTarget = findSubAgentToResume(invocationContext); - if (resumeTarget.isPresent()) { + Optional resumeTargetName = findTransferToResume(invocationContext); + if (resumeTargetName.isPresent()) { + Optional resumeTarget = declaredTransferTarget(resumeTargetName.get()); + // End this agent without running anything if the recorded target is undeclared + // (rejected) or already ended after this agent's latest transfer (e.g. on the + // target's own rejected transfer). + if (resumeTarget.isEmpty() + || endedAfterLatestTransfer(invocationContext, resumeTarget.get())) { + return endOfAgentAndRecord(invocationContext); + } AtomicBoolean resumePaused = new AtomicBoolean(false); return resumeTarget .get() @@ -710,6 +718,7 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { paused.set(true); } }) + .onErrorResumeNext(error -> endIfOwnRejectedTransfer(invocationContext, error)) .concatWith( Flowable.defer( () -> { @@ -724,6 +733,20 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { }); } + /** + * Emits an end-of-agent marker before {@code error} when it is a {@link + * TransferRejectedException} from this agent, so a resume does not re-run it. Any other error, + * including a sub-agent's rejection, passes through and leaves this agent open. + */ + private Flowable endIfOwnRejectedTransfer( + InvocationContext invocationContext, Throwable error) { + if (error instanceof TransferRejectedException rejected + && rejected.agentName().equals(name())) { + return endOfAgentAndRecord(invocationContext).concatWith(Flowable.error(error)); + } + return Flowable.error(error); + } + /** * Returns the name of the agent this agent transferred to in {@code event}, or empty when {@code * event} is not such a transfer. A {@code transfer_to_agent} function response is required, as in @@ -740,33 +763,48 @@ private Optional transferTargetNameFrom(Event event) { } /** - * 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. + * 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 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. - for (BaseAgent target : transferTargets()) { - if (target.name().equals(targetName)) { - return target; + 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 rootAgent() - .findAgent(targetName) - .orElseThrow( - () -> - new IllegalStateException( - "Cannot resume agent " - + name() - + ": the agent it transferred to is no longer in the agent tree.")); + return false; } /** - * When this agent is being resumed, returns the sub-agent it had transferred to (so the resume - * continues that sub-agent), or empty when this agent should continue itself. + * Returns whether {@code event} is {@code author}'s end-of-agent marker, which has no content. */ - private Optional findSubAgentToResume(InvocationContext context) { + private static boolean isEndMarker(Event event, String author) { + return author.equals(event.author()) + && event.actions().endOfAgent() + && event.content().isEmpty(); + } + + /** + * Returns the declared target named {@code targetName}, or empty when it is not one, so a + * same-named agent elsewhere in the tree is never resumed in its place. + */ + private Optional declaredTransferTarget(String targetName) { + return transferTargets().stream() + .filter(target -> target.name().equals(targetName)) + .findFirst(); + } + + /** + * When this agent is being resumed, returns the name of the agent it had transferred to, or empty + * when this agent should continue itself. The name need not be a declared target. + */ + private Optional findTransferToResume(InvocationContext context) { ImmutableList events = context.events(/* currentInvocation= */ true, /* currentBranch= */ true); if (events.isEmpty()) { @@ -774,7 +812,7 @@ private Optional findSubAgentToResume(InvocationContext context) { } Event lastEvent = Iterables.getLast(events); if (name().equals(lastEvent.author())) { - return transferTargetNameFrom(lastEvent).map(this::resolveTransferTarget); + return transferTargetNameFrom(lastEvent); } if (Objects.equals(lastEvent.author(), Role.USER) && !lastEvent.functionResponses().isEmpty()) { // IAE (not ISE): an unresolvable resume surfaces through Runner.runAsync's IAE contract. @@ -793,9 +831,14 @@ private Optional 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 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 targetName = transferTargetNameFrom(event); if (targetName.isPresent()) { - return Optional.of(resolveTransferTarget(targetName.get())); + return targetName; } } return Optional.empty(); diff --git a/core/src/main/java/com/google/adk/agents/TransferRejectedException.java b/core/src/main/java/com/google/adk/agents/TransferRejectedException.java new file mode 100644 index 000000000..51adc2dc5 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/TransferRejectedException.java @@ -0,0 +1,38 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.agents; + +/** + * 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. + */ +public final class TransferRejectedException extends IllegalStateException { + + private final String agentName; + + /** Creates an exception for a transfer rejected from {@code agentName} (the calling agent). */ + public TransferRejectedException(String agentName, String message) { + super(message); + this.agentName = agentName; + } + + /** Returns the name of the agent whose transfer was rejected. */ + public String agentName() { + return agentName; + } +} 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 b953f0a60..e53283e54 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 @@ -27,6 +27,7 @@ import com.google.adk.agents.LlmAgent; import com.google.adk.agents.ReadonlyContext; import com.google.adk.agents.RunConfig.StreamingMode; +import com.google.adk.agents.TransferRejectedException; import com.google.adk.events.Event; import com.google.adk.flows.BaseFlow; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; @@ -424,9 +425,9 @@ private Single handleAfterModelCallback( * *

Handles early termination, LLM call limits, and agent transfer if needed. * - * @return A {@link Flowable} of {@link Event} objects from this step. + * @return A {@link Flowable} of {@link Event} objects from this step, or an error with {@link + * TransferRejectedException} if a transfer names an agent this agent may not transfer to. * @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations. - * @throws IllegalStateException if a transfer agent is specified but not found. */ private Flowable runOneStep(Context spanContext, InvocationContext context) { AtomicReference llmRequestRef = new AtomicReference<>(LlmRequest.builder().build()); @@ -802,34 +803,17 @@ private Flowable followTransfer( return self; } String agentToTransfer = event.actions().transferToAgent().get(); - Optional 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 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 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)))); } @@ -840,13 +824,34 @@ private Flowable followTransfer( */ private static Optional 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 TransferRejectedException transferRejected(LlmAgent from, String name) { + if (name.equals(from.name())) { + return new TransferRejectedException( + from.name(), "Agent '" + from.name() + "' cannot transfer to itself."); + } + if (from.rootAgent().findAgent(name).isEmpty()) { + return new TransferRejectedException( + from.name(), + "Agent '" + + from.name() + + "' requested a transfer to an agent that is not in the agent tree."); + } + return new TransferRejectedException( + from.name(), + "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 diff --git a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java index d017860ce..30c6d9627 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java @@ -16,6 +16,7 @@ package com.google.adk.agents; +import static com.google.adk.testing.TestUtils.END_OF_AGENT; import static com.google.adk.testing.TestUtils.assertEqualIgnoringFunctionIds; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; @@ -25,6 +26,7 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; @@ -645,10 +647,12 @@ public void run_withExampleTool_doesNotAddFunctionDeclarations() { // ---- Resumability: resume into a transferred sub-agent (ResumableLlmAgentTest parity). ---- private static InvocationContext resumableContextWithSeededEvent( - LlmAgent rootAgent, Event seededEvent) { + LlmAgent rootAgent, Event... seededEvents) { InMemorySessionService sessionService = new InMemorySessionService(); Session session = sessionService.createSession("app", "user").blockingGet(); - var unused = sessionService.appendEvent(session, seededEvent).blockingGet(); + for (Event seededEvent : seededEvents) { + var unused = sessionService.appendEvent(session, seededEvent).blockingGet(); + } return InvocationContext.builder() .sessionService(sessionService) .artifactService(new InMemoryArtifactService()) @@ -687,6 +691,24 @@ private static Event transferEvent(String author, String targetAgent) { .build(); } + private static Event endOfAgentEvent(String id, String author) { + return Event.builder() + .id(id) + .invocationId("inv") + .author(author) + .actions(EventActions.builder().endOfAgent(true).build()) + .build(); + } + + private static Event userEvent(String id, String text) { + return Event.builder() + .id(id) + .invocationId("inv") + .author("user") + .content(Content.fromParts(Part.fromText(text))) + .build(); + } + @Test public void runAsync_resumeFromTransferCall_runsTransferredSubAgent() { LlmAgent sub = @@ -835,6 +857,78 @@ 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() { + TestLlm subLlm = createTestLlm(createTextLlmResponse("sub should not run")); + LlmAgent sub = createTestAgentBuilder(subLlm).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"), + endOfAgentEvent("e1", "root"), + userEvent("u1", "next")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyEvents(events)).contains("root: root runs again"); + assertThat(subLlm.getRequests()).isEmpty(); + } + + // 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"), + endOfAgentEvent("e1", "sub"), + endOfAgentEvent("e2", "root"), + transferEvent("root", "sub").toBuilder().id("t2").build()); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true); + + List 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"), + endOfAgentEvent("e1", "root").toBuilder() + .content(Content.fromParts(Part.fromText("ending"))) + .build(), + userEvent("u1", "next")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List 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 @@ -874,17 +968,42 @@ public void transferTargets_noLlmAgentParent_returnsOnlySubAgents() { assertThat(subject.transferTargets().stream().map(BaseAgent::name)).containsExactly("child"); } + // The transferred-to agent is gone: the resume treats the transfer as rejected. @Test - public void runAsync_resumeFromTransferToMissingAgent_throws() { - // The transferred-to agent is gone: a resumed transfer must fail, as a live one does. + public void runAsync_resumeFromTransferToMissingAgent_endsWithoutRunning() { TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); LlmAgent root = createTestAgentBuilder(rootLlm).name("root").build(); InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "removed_agent")); context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); - var resumed = root.runAsync(context).toList(); - assertThrows(IllegalStateException.class, resumed::blockingGet); + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyResumableEvents(events)).containsExactly("root: " + END_OF_AGENT); + assertThat(rootLlm.getRequests()).isEmpty(); + } + + // A transfer to an undeclared agent recorded with no end marker, as in a session written before + // targets were checked, ends this agent instead of running that agent. + @Test + public void runAsync_resumeFromTransferToUndeclaredAgent_endsWithoutRunningIt() { + TestLlm grandchildLlm = createTestLlm(createTextLlmResponse("grandchild should not run")); + LlmAgent grandchild = createTestAgentBuilder(grandchildLlm).name("grandchild").build(); + LlmAgent sub = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub should not run"))) + .name("sub") + .subAgents(grandchild) + .build(); + TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run")); + LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build(); + InvocationContext context = + resumableContextWithSeededEvent(root, transferEvent("root", "grandchild")); + context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false); + + List events = root.runAsync(context).toList().blockingGet(); + + assertThat(simplifyResumableEvents(events)).containsExactly("root: " + END_OF_AGENT); + assertThat(grandchildLlm.getRequests()).isEmpty(); assertThat(rootLlm.getRequests()).isEmpty(); } 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 fa7abf0a9..702430ff2 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 @@ -16,13 +16,22 @@ package com.google.adk.flows.llmflows; +import static com.google.adk.testing.ResumabilityTestUtils.newSession; +import static com.google.adk.testing.ResumabilityTestUtils.reloadSession; +import static com.google.adk.testing.ResumabilityTestUtils.resumableRunner; +import static com.google.adk.testing.ResumabilityTestUtils.resumeById; +import static com.google.adk.testing.ResumabilityTestUtils.runTurn; +import static com.google.adk.testing.ResumabilityTestUtils.textAgent; +import static com.google.adk.testing.TestUtils.END_OF_AGENT; import static com.google.adk.testing.TestUtils.createInvocationContext; import static com.google.adk.testing.TestUtils.createLlmResponse; import static com.google.adk.testing.TestUtils.createTestAgentBuilder; import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.adk.testing.TestUtils.createTextLlmResponse; import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.adk.testing.TestUtils.simplifyResumableEvents; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; @@ -30,9 +39,12 @@ import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.LoopAgent; +import com.google.adk.agents.ParallelAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; +import com.google.adk.agents.TransferRejectedException; import com.google.adk.events.Event; +import com.google.adk.models.LlmResponse; import com.google.adk.runner.InMemoryRunner; import com.google.adk.runner.Runner; import com.google.adk.sessions.Session; @@ -51,6 +63,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -88,6 +101,31 @@ public Single> runAsync(Map args, ToolContex } } + /** A tool that always fails, counting its calls. */ + private static class FailingTool extends BaseTool { + final AtomicInteger calls = new AtomicInteger(); + + FailingTool() { + super("fail", "Always fails."); + } + + @Override + public Optional declaration() { + return Optional.of( + FunctionDeclaration.builder() + .name(name()) + .description(description()) + .parameters(Schema.builder().type("OBJECT").build()) + .build()); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + calls.incrementAndGet(); + return Single.error(new IllegalStateException("tool failed")); + } + } + // Order is load-bearing: this list is rendered into the model's transfer instruction. @Test public void transferTargets_ordersSubAgentsThenParentThenPeers() { @@ -519,6 +557,381 @@ public void testAutoToLoop() { assertThat(simplifyEvents(actualEvents)).containsExactly("root_agent: response5"); } + // A grandchild is not a sub-agent, parent or peer, so it is not a transfer target. + @Test + public void runAsync_transferToGrandchild_failsAndGrandchildNeverRuns() { + TestLlm grandchildLlm = TestLlm.create(createTextLlmResponse("grandchild ran")); + Runner runner = new InMemoryRunner(rootTransferringToGrandchild(grandchildLlm), "test"); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'root' is not allowed to transfer to the requested agent."); + assertThat(grandchildLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumable_transferToGrandchild_failsAndGrandchildNeverRuns() { + TestLlm grandchildLlm = TestLlm.create(createTextLlmResponse("grandchild ran")); + Runner runner = resumableRunner(rootTransferringToGrandchild(grandchildLlm)); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'root' is not allowed to transfer to the requested agent."); + assertThat(grandchildLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_transferToSelf_failsAndTheAgentDoesNotRunAgain() { + TestLlm rootLlm = TestLlm.create(transferTo("root")); + Runner runner = new InMemoryRunner(root(rootLlm, textAgent("sub")), "test"); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error).hasMessageThat().isEqualTo("Agent 'root' cannot transfer to itself."); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_resumable_transferToSelf_failsAndTheAgentDoesNotRunAgain() { + TestLlm rootLlm = TestLlm.create(transferTo("root")); + Runner runner = resumableRunner(root(rootLlm, textAgent("sub"))); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error).hasMessageThat().isEqualTo("Agent 'root' cannot transfer to itself."); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_transferToParent_withDisallowTransferToParent_fails() { + TestLlm rootLlm = TestLlm.create(transferTo("sub")); + Runner runner = new InMemoryRunner(rootWithSubDisallowingParent(rootLlm), "test"); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent."); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_resumable_transferToParent_withDisallowTransferToParent_fails() { + TestLlm rootLlm = TestLlm.create(transferTo("sub")); + Runner runner = resumableRunner(rootWithSubDisallowingParent(rootLlm)); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent."); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void runAsync_transferToPeer_withDisallowTransferToPeers_failsAndPeerNeverRuns() { + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + Runner runner = + new InMemoryRunner( + rootWithSubDisallowingPeers(TestLlm.create(transferTo("peer")), peerLlm), "test"); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent."); + assertThat(error.agentName()).isEqualTo("sub"); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + @Test + public void + runAsync_resumable_transferToPeer_withDisallowTransferToPeers_failsAndPeerNeverRuns() { + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + Runner runner = + resumableRunner(rootWithSubDisallowingPeers(TestLlm.create(transferTo("peer")), peerLlm)); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent."); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // The model chose the name, so the error names only the calling agent. + @Test + public void runAsync_transferToUnknownAgent_failsAsNotInTree() { + Runner runner = + new InMemoryRunner(root(TestLlm.create(transferTo("missing")), textAgent("sub")), "test"); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'root' requested a transfer to an agent that is not in the agent tree."); + } + + @Test + public void runAsync_resumable_transferToUnknownAgent_failsAsNotInTree() { + Runner runner = resumableRunner(root(TestLlm.create(transferTo("missing")), textAgent("sub"))); + Session session = newSession(runner); + + TransferRejectedException error = + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'root' requested a transfer to an agent that is not in the agent tree."); + } + + @Test + public void runAsync_transferToDuplicateName_runsTheDeclaredTarget() { + TestLlm declaredLlm = TestLlm.create(createTextLlmResponse("declared ran")); + TestLlm namesakeLlm = TestLlm.create(createTextLlmResponse("namesake ran")); + Runner runner = new InMemoryRunner(rootWithNamesakeTargets(declaredLlm, namesakeLlm), "test"); + + runTurn(runner, newSession(runner), "hi"); + + assertThat(declaredLlm.getRequests()).hasSize(1); + assertThat(namesakeLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumable_transferToDuplicateName_runsTheDeclaredTarget() { + TestLlm declaredLlm = TestLlm.create(createTextLlmResponse("declared ran")); + TestLlm namesakeLlm = TestLlm.create(createTextLlmResponse("namesake ran")); + Runner runner = resumableRunner(rootWithNamesakeTargets(declaredLlm, namesakeLlm)); + + runTurn(runner, newSession(runner), "hi"); + + assertThat(declaredLlm.getRequests()).hasSize(1); + assertThat(namesakeLlm.getRequests()).isEmpty(); + } + + // After a rejected transfer, the next turn runs the agent that tried it. + @Test + public void runAsync_afterRejectedPeerTransfer_nextTurnRunsTheTransferringAgent() { + TestLlm subLlm = TestLlm.create(transferTo("peer"), createTextLlmResponse("sub turn 2")); + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + Runner runner = new InMemoryRunner(rootWithSubDisallowingPeers(subLlm, peerLlm), "test"); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "turn 1")); + + ImmutableList turn2 = runTurn(runner, session, "turn 2"); + + assertThat(simplifyEvents(turn2)).containsExactly("sub: sub turn 2"); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + @Test + public void runAsync_resumable_afterRejectedPeerTransfer_nextTurnRunsTheTransferringAgent() { + TestLlm subLlm = TestLlm.create(transferTo("peer"), createTextLlmResponse("sub turn 2")); + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + Runner runner = resumableRunner(rootWithSubDisallowingPeers(subLlm, peerLlm)); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "turn 1")); + + ImmutableList turn2 = runTurn(runner, session, "turn 2"); + + assertThat(simplifyResumableEvents(turn2)) + .containsExactly("sub: sub turn 2", "sub: " + END_OF_AGENT) + .inOrder(); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // A rejected transfer ends the agent that made it, so resuming the invocation runs nothing. + @Test + public void resume_afterRejectedSelfTransfer_runsNothing() { + TestLlm rootLlm = TestLlm.create(transferTo("root")); + Runner runner = resumableRunner(root(rootLlm, textAgent("sub"))); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + String invocationId = reloadSession(runner, session).immutableEvents().get(0).invocationId(); + + ImmutableList resumed = resumeById(runner, session, invocationId); + + assertThat(resumed).isEmpty(); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + @Test + public void resume_afterTransferToUnknownAgent_runsNothing() { + TestLlm rootLlm = TestLlm.create(transferTo("missing")); + Runner runner = resumableRunner(root(rootLlm, textAgent("sub"))); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + String invocationId = reloadSession(runner, session).immutableEvents().get(0).invocationId(); + + ImmutableList resumed = resumeById(runner, session, invocationId); + + assertThat(resumed).isEmpty(); + assertThat(rootLlm.getRequests()).hasSize(1); + } + + // The transferring sub-agent ends before the run fails, so a resume only closes the root. + @Test + public void resume_afterRejectedPeerTransfer_runsNeitherThePeerNorTheTransferringAgent() { + TestLlm subLlm = TestLlm.create(transferTo("peer")); + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + Runner runner = resumableRunner(rootWithSubDisallowingPeers(subLlm, peerLlm)); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + ImmutableList firstTurnEvents = reloadSession(runner, session).immutableEvents(); + ImmutableList firstTurn = simplifyResumableEvents(firstTurnEvents); + String invocationId = firstTurnEvents.get(0).invocationId(); + + ImmutableList resumed = resumeById(runner, session, invocationId); + + assertThat(firstTurn.subList(firstTurn.size() - 2, firstTurn.size())) + .containsExactly( + "sub: FunctionResponse(name=transfer_to_agent, response={})", "sub: " + END_OF_AGENT) + .inOrder(); + assertThat(simplifyResumableEvents(resumed)).containsExactly("root: " + END_OF_AGENT); + assertThat(subLlm.getRequests()).hasSize(1); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // Under a ParallelAgent, sub writes its end marker on a descendant branch of root's. + @Test + public void resume_underParallelAgent_afterRejectedPeerTransfer_onlyClosesTheAncestors() { + TestLlm subLlm = TestLlm.create(transferTo("peer")); + TestLlm peerLlm = TestLlm.create(ImmutableList.of()); + ParallelAgent par = + ParallelAgent.builder() + .name("par") + .subAgents(ImmutableList.of(rootWithSubDisallowingPeers(subLlm, peerLlm))) + .build(); + Runner runner = resumableRunner(par); + Session session = newSession(runner); + assertThrows(TransferRejectedException.class, () -> runTurn(runner, session, "hi")); + String invocationId = reloadSession(runner, session).immutableEvents().get(0).invocationId(); + + ImmutableList resumed = resumeById(runner, session, invocationId); + + assertThat(simplifyResumableEvents(resumed)) + .containsExactly("root: " + END_OF_AGENT, "par: " + END_OF_AGENT) + .inOrder(); + assertThat(subLlm.getRequests()).hasSize(1); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + // Only a rejected transfer ends the agent; after any other failure a resume runs it again. + @Test + public void resume_afterModelError_runsTheAgentAgain() { + TestLlm rootLlm = + TestLlm.create(/* responses= */ null, new IllegalStateException("model failed")); + Runner runner = resumableRunner(root(rootLlm, textAgent("sub"))); + Session session = newSession(runner); + assertThrows(IllegalStateException.class, () -> runTurn(runner, session, "hi")); + ImmutableList firstTurnEvents = reloadSession(runner, session).immutableEvents(); + String invocationId = firstTurnEvents.get(0).invocationId(); + + assertThrows(IllegalStateException.class, () -> resumeById(runner, session, invocationId)); + + assertThat(simplifyResumableEvents(firstTurnEvents)).doesNotContain("root: " + END_OF_AGENT); + assertThat(rootLlm.getRequests()).hasSize(2); + } + + @Test + public void resume_afterToolError_runsTheAgentAgain() { + TestLlm rootLlm = + TestLlm.create( + createLlmResponse(Content.fromParts(Part.fromFunctionCall("fail", ImmutableMap.of())))); + FailingTool tool = new FailingTool(); + LlmAgent root = + createTestAgentBuilder(rootLlm) + .name("root") + .tools(tool) + .subAgents(textAgent("sub")) + .build(); + Runner runner = resumableRunner(root); + Session session = newSession(runner); + assertThrows(IllegalStateException.class, () -> runTurn(runner, session, "hi")); + ImmutableList firstTurnEvents = reloadSession(runner, session).immutableEvents(); + String invocationId = firstTurnEvents.get(0).invocationId(); + + // The resume replays the pending call, so the tool fails again. + assertThrows(IllegalStateException.class, () -> resumeById(runner, session, invocationId)); + + assertThat(simplifyResumableEvents(firstTurnEvents)).doesNotContain("root: " + END_OF_AGENT); + assertThat(tool.calls.get()).isEqualTo(2); + } + + private static LlmResponse transferTo(String agentName) { + return createLlmResponse(Content.fromParts(createTransferCallPart(agentName))); + } + + private static LlmAgent root(TestLlm llm, BaseAgent... subAgents) { + return createTestAgentBuilder(llm).name("root").subAgents(subAgents).build(); + } + + /** A root whose model transfers to its grandchild, which runs on {@code grandchildLlm}. */ + private static LlmAgent rootTransferringToGrandchild(TestLlm grandchildLlm) { + LlmAgent grandchild = createTestAgentBuilder(grandchildLlm).name("grandchild").build(); + LlmAgent sub = + createTestAgentBuilder(TestLlm.create(ImmutableList.of())) + .name("sub") + .subAgents(grandchild) + .build(); + return root(TestLlm.create(transferTo("grandchild")), sub); + } + + /** A root, on {@code rootLlm}, whose sub-agent transfers back despite disallowing it. */ + private static LlmAgent rootWithSubDisallowingParent(TestLlm rootLlm) { + LlmAgent sub = + createTestAgentBuilder(TestLlm.create(transferTo("root"))) + .name("sub") + .disallowTransferToParent(true) + .build(); + // Gives sub a target, so it is offered the transfer tool. + return root(rootLlm, sub, textAgent("peer")); + } + + /** A root that transfers to {@code sub}, on {@code subLlm}, which disallows peer transfers. */ + private static LlmAgent rootWithSubDisallowingPeers(TestLlm subLlm, TestLlm peerLlm) { + LlmAgent sub = createTestAgentBuilder(subLlm).name("sub").disallowTransferToPeers(true).build(); + LlmAgent peer = createTestAgentBuilder(peerLlm).name("peer").build(); + return root(TestLlm.create(transferTo("sub")), sub, peer); + } + + /** + * A root whose sub-agent "shared" runs on {@code declaredLlm}, a deeper one on {@code + * namesakeLlm}. + */ + private static LlmAgent rootWithNamesakeTargets(TestLlm declaredLlm, TestLlm namesakeLlm) { + LlmAgent otherBranch = + createTestAgentBuilder(TestLlm.create(ImmutableList.of())) + .name("other_branch") + .subAgents(createTestAgentBuilder(namesakeLlm).name("shared").build()) + .build(); + LlmAgent declared = createTestAgentBuilder(declaredLlm).name("shared").build(); + // other_branch comes first, so a depth-first search of the tree reaches the namesake first. + return root(TestLlm.create(transferTo("shared")), otherBranch, declared); + } + private Runner getRunnerAndCreateSession(LlmAgent agent, Session session) { Runner runner = new InMemoryRunner(agent, session.appName()); // Ensure the session exists before running the agent.