From 1ec34b9e8ba5aef8a27b3df13c30d9dd54ed3000 Mon Sep 17 00:00:00 2001 From: Kamil Tomaszek Date: Thu, 15 Oct 2026 03:39:37 -0700 Subject: [PATCH] fix: restrict transfer_to_agent to sub-agents, parent, and peers - Rejected transfers throw `IllegalArgumentException` instead of `IllegalStateException`. PiperOrigin-RevId: 1000022357 --- .../java/com/google/adk/agents/LlmAgent.java | 66 +++- .../adk/flows/llmflows/BaseLlmFlow.java | 49 +-- .../com/google/adk/agents/LlmAgentTest.java | 106 ++++++ .../adk/flows/llmflows/AgentTransferTest.java | 320 ++++++++++++++++++ .../adk/runner/RunnerResumabilityTest.java | 2 +- 5 files changed, 515 insertions(+), 28 deletions(-) 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..b9902f60a 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -679,6 +679,9 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { if (invocationContext.agentStates().containsKey(name())) { Optional 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 runAsyncImpl(InvocationContext invocationContext) { Flowable 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 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 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())); } 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 43faef5cf..d9e02c655 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 @@ -426,7 +426,7 @@ private Single 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 runOneStep(Context spanContext, InvocationContext context) { AtomicReference llmRequestRef = new AtomicReference<>(LlmRequest.builder().build()); @@ -801,34 +801,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)))); } @@ -839,13 +822,31 @@ 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 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 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..e0d323f60 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/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 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 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 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 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..0f2a8e736 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; @@ -33,6 +42,7 @@ import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; 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; @@ -519,6 +529,316 @@ 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); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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", "sub ran")), "test"); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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", "sub ran"))); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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(createTextLlmResponse("peer ran")); + Runner runner = + new InMemoryRunner( + rootWithSubDisallowingPeers(TestLlm.create(transferTo("peer")), peerLlm), "test"); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> runTurn(runner, session, "hi")); + + assertThat(error) + .hasMessageThat() + .isEqualTo("Agent 'sub' is not allowed to transfer to the requested agent."); + assertThat(peerLlm.getRequests()).isEmpty(); + } + + @Test + public void + runAsync_resumable_transferToPeer_withDisallowTransferToPeers_failsAndPeerNeverRuns() { + TestLlm peerLlm = TestLlm.create(createTextLlmResponse("peer ran")); + Runner runner = + resumableRunner(rootWithSubDisallowingPeers(TestLlm.create(transferTo("peer")), peerLlm)); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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", "sub ran")), "test"); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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", "sub ran"))); + Session session = newSession(runner); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.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(createTextLlmResponse("peer ran")); + Runner runner = new InMemoryRunner(rootWithSubDisallowingPeers(subLlm, peerLlm), "test"); + Session session = newSession(runner); + assertThrows(IllegalArgumentException.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(createTextLlmResponse("peer ran")); + Runner runner = resumableRunner(rootWithSubDisallowingPeers(subLlm, peerLlm)); + Session session = newSession(runner); + assertThrows(IllegalArgumentException.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", "sub ran"))); + Session session = newSession(runner); + assertThrows(IllegalArgumentException.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", "sub ran"))); + Session session = newSession(runner); + assertThrows(IllegalArgumentException.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(createTextLlmResponse("peer ran")); + Runner runner = resumableRunner(rootWithSubDisallowingPeers(subLlm, peerLlm)); + Session session = newSession(runner); + assertThrows(IllegalArgumentException.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(); + } + + 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(createTextLlmResponse("sub ran"))) + .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", "peer ran")); + } + + /** 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(createTextLlmResponse("other ran"))) + .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. diff --git a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java index 147d4636d..d442f8df8 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -940,7 +940,7 @@ public void resume_replayedTransferToUndeclaredAgent_fails() { .test() .assertError( error -> - error instanceof IllegalStateException + error instanceof IllegalArgumentException && error.getMessage().contains("is not allowed to transfer")); assertThat(peerLlm.getRequests()).isEmpty(); }