Skip to content

Commit 84bae34

Browse files
Mateusz Krawieccopybara-github
authored andcommitted
revert: restrict turn-based transfer_to_agent to declared targets
Rolled back due to a breakage. PiperOrigin-RevId: 1005598267
1 parent 18554a1 commit 84bae34

5 files changed

Lines changed: 56 additions & 674 deletions

File tree

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

Lines changed: 25 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -678,16 +678,8 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
678678
() -> {
679679
// Resumed after a transfer: continue the sub-agent instead of re-invoking the model.
680680
if (invocationContext.agentStates().containsKey(name())) {
681-
Optional<String> resumeTargetName = findTransferToResume(invocationContext);
682-
if (resumeTargetName.isPresent()) {
683-
Optional<BaseAgent> resumeTarget = declaredTransferTarget(resumeTargetName.get());
684-
// End this agent without running anything if the recorded target is undeclared
685-
// (rejected) or already ended after this agent's latest transfer (e.g. on the
686-
// target's own rejected transfer).
687-
if (resumeTarget.isEmpty()
688-
|| endedAfterLatestTransfer(invocationContext, resumeTarget.get())) {
689-
return endOfAgentAndRecord(invocationContext);
690-
}
681+
Optional<BaseAgent> resumeTarget = findSubAgentToResume(invocationContext);
682+
if (resumeTarget.isPresent()) {
691683
AtomicBoolean resumePaused = new AtomicBoolean(false);
692684
return resumeTarget
693685
.get()
@@ -719,7 +711,6 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
719711
paused.set(true);
720712
}
721713
})
722-
.onErrorResumeNext(error -> endIfOwnRejectedTransfer(invocationContext, error))
723714
.concatWith(
724715
Flowable.defer(
725716
() -> {
@@ -734,20 +725,6 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
734725
});
735726
}
736727

737-
/**
738-
* Emits an end-of-agent marker before {@code error} when it is a {@link
739-
* TransferRejectedException} from this agent, so a resume does not re-run it. Any other error,
740-
* including a sub-agent's rejection, passes through and leaves this agent open.
741-
*/
742-
private Flowable<Event> endIfOwnRejectedTransfer(
743-
InvocationContext invocationContext, Throwable error) {
744-
if (error instanceof TransferRejectedException rejected
745-
&& rejected.agentName().equals(name())) {
746-
return endOfAgentAndRecord(invocationContext).concatWith(Flowable.error(error));
747-
}
748-
return Flowable.error(error);
749-
}
750-
751728
/**
752729
* Returns the name of the agent this agent transferred to in {@code event}, or empty when {@code
753730
* event} is not such a transfer. A {@code transfer_to_agent} function response is required, as in
@@ -764,56 +741,41 @@ private Optional<String> transferTargetNameFrom(Event event) {
764741
}
765742

766743
/**
767-
* Returns whether {@code target} ended after this agent's latest transfer, as it does on its own
768-
* rejected transfer. An end marker from before that transfer, such as a past loop iteration, does
769-
* not count.
744+
* Resolves a transfer target recorded in history back to its agent. A live transfer fails in
745+
* {@code BaseLlmFlow} when the target is gone; a resumed one has no such check, and returning
746+
* empty here would silently re-run this agent instead of the sub-agent it handed off to.
770747
*/
771-
private boolean endedAfterLatestTransfer(InvocationContext context, BaseAgent target) {
772-
// All branches: under a ParallelAgent, the target writes its marker on a descendant branch.
773-
for (Event event :
774-
context.events(/* currentInvocation= */ true, /* currentBranch= */ false).reverse()) {
775-
if (transferTargetNameFrom(event).isPresent()) {
776-
return false;
777-
}
778-
if (isEndMarker(event, target.name())) {
779-
return true;
748+
private BaseAgent resolveTransferTarget(String targetName) {
749+
// Declared targets first, as in Python, so a same-named agent on an unrelated branch is not
750+
// resumed into; then the same tree-wide search the live transfer path uses.
751+
for (BaseAgent target : transferTargets()) {
752+
if (target.name().equals(targetName)) {
753+
return target;
780754
}
781755
}
782-
return false;
756+
return rootAgent()
757+
.findAgent(targetName)
758+
.orElseThrow(
759+
() ->
760+
new IllegalStateException(
761+
"Cannot resume agent "
762+
+ name()
763+
+ ": the agent it transferred to is no longer in the agent tree."));
783764
}
784765

785766
/**
786-
* Returns whether {@code event} is {@code author}'s end-of-agent marker, which has no content.
767+
* When this agent is being resumed, returns the sub-agent it had transferred to (so the resume
768+
* continues that sub-agent), or empty when this agent should continue itself.
787769
*/
788-
private static boolean isEndMarker(Event event, String author) {
789-
return author.equals(event.author())
790-
&& event.actions().endOfAgent()
791-
&& event.content().isEmpty();
792-
}
793-
794-
/**
795-
* Returns the declared target named {@code targetName}, or empty when it is not one, so a
796-
* same-named agent elsewhere in the tree is never resumed in its place.
797-
*/
798-
private Optional<BaseAgent> declaredTransferTarget(String targetName) {
799-
return transferTargets().stream()
800-
.filter(target -> target.name().equals(targetName))
801-
.findFirst();
802-
}
803-
804-
/**
805-
* When this agent is being resumed, returns the name of the agent it had transferred to, or empty
806-
* when this agent should continue itself. The name need not be a declared target.
807-
*/
808-
private Optional<String> findTransferToResume(InvocationContext context) {
770+
private Optional<BaseAgent> findSubAgentToResume(InvocationContext context) {
809771
ImmutableList<Event> events =
810772
context.events(/* currentInvocation= */ true, /* currentBranch= */ true);
811773
if (events.isEmpty()) {
812774
return Optional.empty();
813775
}
814776
Event lastEvent = Iterables.getLast(events);
815777
if (name().equals(lastEvent.author())) {
816-
return transferTargetNameFrom(lastEvent);
778+
return transferTargetNameFrom(lastEvent).map(this::resolveTransferTarget);
817779
}
818780
if (Objects.equals(lastEvent.author(), Role.USER) && !lastEvent.functionResponses().isEmpty()) {
819781
// IAE (not ISE): an unresolvable resume surfaces through Runner.runAsync's IAE contract.
@@ -832,14 +794,9 @@ private Optional<String> findTransferToResume(InvocationContext context) {
832794
}
833795
// A user event answering no call falls through; Python fails the resume instead.
834796
for (int i = events.size() - 2; i >= 0; i--) {
835-
Event event = events.get(i);
836-
// Transfers before this agent's own end marker belong to a finished run, such as a past loop.
837-
if (isEndMarker(event, name())) {
838-
break;
839-
}
840-
Optional<String> targetName = transferTargetNameFrom(event);
797+
Optional<String> targetName = transferTargetNameFrom(events.get(i));
841798
if (targetName.isPresent()) {
842-
return targetName;
799+
return Optional.of(resolveTransferTarget(targetName.get()));
843800
}
844801
}
845802
return Optional.empty();

‎core/src/main/java/com/google/adk/agents/TransferRejectedException.java‎

Lines changed: 0 additions & 38 deletions
This file was deleted.

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

Lines changed: 25 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import com.google.adk.agents.LlmAgent;
2828
import com.google.adk.agents.ReadonlyContext;
2929
import com.google.adk.agents.RunConfig.StreamingMode;
30-
import com.google.adk.agents.TransferRejectedException;
3130
import com.google.adk.events.Event;
3231
import com.google.adk.flows.BaseFlow;
3332
import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult;
@@ -425,9 +424,9 @@ private Single<LlmResponse> handleAfterModelCallback(
425424
*
426425
* <p>Handles early termination, LLM call limits, and agent transfer if needed.
427426
*
428-
* @return A {@link Flowable} of {@link Event} objects from this step, or an error with {@link
429-
* TransferRejectedException} if a transfer names an agent this agent may not transfer to.
427+
* @return A {@link Flowable} of {@link Event} objects from this step.
430428
* @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
429+
* @throws IllegalStateException if a transfer agent is specified but not found.
431430
*/
432431
private Flowable<Event> runOneStep(Context spanContext, InvocationContext context) {
433432
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
@@ -803,17 +802,34 @@ private Flowable<Event> followTransfer(
803802
return self;
804803
}
805804
String agentToTransfer = event.actions().transferToAgent().get();
806-
LlmAgent agent = (LlmAgent) context.agent();
807-
Optional<BaseAgent> target = declaredTransferTarget(agent, agentToTransfer);
808-
if (target.isEmpty()) {
809-
return self.concatWith(Flowable.error(transferRejected(agent, agentToTransfer)));
805+
Optional<BaseAgent> nextAgent = context.agent().rootAgent().findAgent(agentToTransfer);
806+
if (nextAgent.isEmpty()) {
807+
logger.error("Agent not found: {}", agentToTransfer);
808+
return self.concatWith(
809+
Flowable.error(new IllegalStateException("Agent not found: " + agentToTransfer)));
810+
}
811+
BaseAgent target;
812+
// Resumable only: restricting the default flow too would move behavior for existing callers.
813+
if (context.isResumable() && context.agent() instanceof LlmAgent llmAgent) {
814+
Optional<BaseAgent> declared = declaredTransferTarget(llmAgent, agentToTransfer);
815+
if (declared.isEmpty()) {
816+
return self.concatWith(
817+
Flowable.error(
818+
new IllegalStateException(
819+
"Agent "
820+
+ llmAgent.name()
821+
+ " is not allowed to transfer to the requested agent.")));
822+
}
823+
target = declared.get();
824+
} else {
825+
target = nextAgent.get();
810826
}
811827
// Otherwise the target can build its request from a session missing the transfer event.
812828
return self.concatWith(
813829
PersistBarrier.awaitPersisted(context, ImmutableList.of(event))
814830
.andThen(
815831
// Not around the barrier: andThen subscribes the target on the append's thread.
816-
Flowable.defer(() -> target.get().runAsync(context))
832+
Flowable.defer(() -> target.runAsync(context))
817833
.compose(Tracing.withContext(spanContext))));
818834
}
819835

@@ -824,34 +840,13 @@ private Flowable<Event> followTransfer(
824840
*/
825841
private static Optional<BaseAgent> declaredTransferTarget(LlmAgent from, String name) {
826842
if (name.equals(from.name())) {
827-
return Optional.empty();
843+
return Optional.of(from);
828844
}
829845
return AgentTransfer.transferTargets(from).stream()
830846
.filter(candidate -> candidate.name().equals(name))
831847
.findFirst();
832848
}
833849

834-
/**
835-
* Returns the error for a transfer {@link #declaredTransferTarget} rejected. The message names
836-
* only {@code from}, since {@code name} is model output.
837-
*/
838-
private static TransferRejectedException transferRejected(LlmAgent from, String name) {
839-
if (name.equals(from.name())) {
840-
return new TransferRejectedException(
841-
from.name(), "Agent '" + from.name() + "' cannot transfer to itself.");
842-
}
843-
if (from.rootAgent().findAgent(name).isEmpty()) {
844-
return new TransferRejectedException(
845-
from.name(),
846-
"Agent '"
847-
+ from.name()
848-
+ "' requested a transfer to an agent that is not in the agent tree.");
849-
}
850-
return new TransferRejectedException(
851-
from.name(),
852-
"Agent '" + from.name() + "' is not allowed to transfer to the requested agent.");
853-
}
854-
855850
/**
856851
* Runs {@code callEvent}'s function calls and emits what follows from them: the tool-confirmation
857852
* request, the function response, and any structured final response. Shared by the normal

0 commit comments

Comments
 (0)