Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 63 additions & 3 deletions core/src/main/java/com/google/adk/agents/LlmAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,9 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
if (invocationContext.agentStates().containsKey(name())) {
Optional<BaseAgent> resumeTarget = findSubAgentToResume(invocationContext);
if (resumeTarget.isPresent()) {
if (endedAfterLatestTransfer(invocationContext, resumeTarget.get())) {
return endOfAgentAndRecord(invocationContext);
}
AtomicBoolean resumePaused = new AtomicBoolean(false);
return resumeTarget
.get()
Expand All @@ -703,13 +706,21 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
Flowable<Event> events =
llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState);
AtomicBoolean paused = new AtomicBoolean(false);
AtomicBoolean lastEventRejectedTransfer = new AtomicBoolean(false);
return events
.doOnNext(
event -> {
if (invocationContext.shouldPauseInvocation(event)) {
paused.set(true);
}
lastEventRejectedTransfer.set(isRejectedTransfer(event));
})
// End this agent before a rejected transfer fails, so a resume does not rerun it.
.onErrorResumeNext(
error ->
lastEventRejectedTransfer.get()
? endOfAgentAndRecord(invocationContext).concatWith(Flowable.error(error))
: Flowable.error(error))
.concatWith(
Flowable.defer(
() -> {
Expand Down Expand Up @@ -739,14 +750,58 @@ private Optional<String> transferTargetNameFrom(Event event) {
return event.actions().transferToAgent().filter(target -> !target.equals(name()));
}

/**
* Returns whether {@code event} is a transfer by this agent that {@code BaseLlmFlow} rejects: one
* to itself or to an agent outside {@link #transferTargets()}.
*/
private boolean isRejectedTransfer(Event event) {
return name().equals(event.author())
&& event
.actions()
.transferToAgent()
.filter(
target ->
target.equals(name())
|| transferTargets().stream()
.noneMatch(agent -> agent.name().equals(target)))
.isPresent();
}

/**
* Returns whether {@code target} ended after this agent's latest transfer, as it does on its own
* rejected transfer. An end marker from before that transfer, such as a past loop iteration, does
* not count.
*/
private boolean endedAfterLatestTransfer(InvocationContext context, BaseAgent target) {
// All branches: under a ParallelAgent, the target writes its marker on a descendant branch.
for (Event event :
context.events(/* currentInvocation= */ true, /* currentBranch= */ false).reverse()) {
if (transferTargetNameFrom(event).isPresent()) {
return false;
}
if (isEndMarker(event, target.name())) {
return true;
}
}
return false;
}

/**
* Returns whether {@code event} is {@code author}'s end-of-agent marker, which has no content.
*/
private static boolean isEndMarker(Event event, String author) {
return author.equals(event.author())
&& event.actions().endOfAgent()
&& event.content().isEmpty();
}

/**
* Resolves a transfer target recorded in history back to its agent. A live transfer fails in
* {@code BaseLlmFlow} when the target is gone; a resumed one has no such check, and returning
* empty here would silently re-run this agent instead of the sub-agent it handed off to.
*/
private BaseAgent resolveTransferTarget(String targetName) {
// Declared targets first, as in Python, so a same-named agent on an unrelated branch is not
// resumed into; then the same tree-wide search the live transfer path uses.
// Check declared targets first, so a same-named agent elsewhere is not resumed instead.
for (BaseAgent target : transferTargets()) {
if (target.name().equals(targetName)) {
return target;
Expand Down Expand Up @@ -793,7 +848,12 @@ private Optional<BaseAgent> findSubAgentToResume(InvocationContext context) {
}
// A user event answering no call falls through; Python fails the resume instead.
for (int i = events.size() - 2; i >= 0; i--) {
Optional<String> targetName = transferTargetNameFrom(events.get(i));
Event event = events.get(i);
// Transfers before this agent's own end marker belong to a finished run, such as a past loop.
if (isEndMarker(event, name())) {
break;
}
Optional<String> targetName = transferTargetNameFrom(event);
if (targetName.isPresent()) {
return Optional.of(resolveTransferTarget(targetName.get()));
}
Expand Down
49 changes: 25 additions & 24 deletions core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,7 +426,7 @@ private Single<LlmResponse> handleAfterModelCallback(
*
* @return A {@link Flowable} of {@link Event} objects from this step.
* @throws LlmCallsLimitExceededException if the agent exceeds allowed LLM invocations.
* @throws IllegalStateException if a transfer agent is specified but not found.
* @throws IllegalArgumentException if a transfer names an agent this agent may not transfer to.
*/
private Flowable<Event> runOneStep(Context spanContext, InvocationContext context) {
AtomicReference<LlmRequest> llmRequestRef = new AtomicReference<>(LlmRequest.builder().build());
Expand Down Expand Up @@ -801,34 +801,17 @@ private Flowable<Event> followTransfer(
return self;
}
String agentToTransfer = event.actions().transferToAgent().get();
Optional<BaseAgent> nextAgent = context.agent().rootAgent().findAgent(agentToTransfer);
if (nextAgent.isEmpty()) {
logger.error("Agent not found: {}", agentToTransfer);
return self.concatWith(
Flowable.error(new IllegalStateException("Agent not found: " + agentToTransfer)));
}
BaseAgent target;
// Resumable only: restricting the default flow too would move behavior for existing callers.
if (context.isResumable() && context.agent() instanceof LlmAgent llmAgent) {
Optional<BaseAgent> declared = declaredTransferTarget(llmAgent, agentToTransfer);
if (declared.isEmpty()) {
return self.concatWith(
Flowable.error(
new IllegalStateException(
"Agent "
+ llmAgent.name()
+ " is not allowed to transfer to the requested agent.")));
}
target = declared.get();
} else {
target = nextAgent.get();
LlmAgent agent = (LlmAgent) context.agent();
Optional<BaseAgent> target = declaredTransferTarget(agent, agentToTransfer);
if (target.isEmpty()) {
return self.concatWith(Flowable.error(transferRejected(agent, agentToTransfer)));
}
// Otherwise the target can build its request from a session missing the transfer event.
return self.concatWith(
PersistBarrier.awaitPersisted(context, ImmutableList.of(event))
.andThen(
// Not around the barrier: andThen subscribes the target on the append's thread.
Flowable.defer(() -> target.runAsync(context))
Flowable.defer(() -> target.get().runAsync(context))
.compose(Tracing.withContext(spanContext))));
}

Expand All @@ -839,13 +822,31 @@ private Flowable<Event> followTransfer(
*/
private static Optional<BaseAgent> declaredTransferTarget(LlmAgent from, String name) {
if (name.equals(from.name())) {
return Optional.of(from);
return Optional.empty();
}
return AgentTransfer.transferTargets(from).stream()
.filter(candidate -> candidate.name().equals(name))
.findFirst();
}

/**
* Returns the error for a transfer {@link #declaredTransferTarget} rejected. The message names
* only {@code from}, since {@code name} is model output.
*/
private static IllegalArgumentException transferRejected(LlmAgent from, String name) {
if (name.equals(from.name())) {
return new IllegalArgumentException("Agent '" + from.name() + "' cannot transfer to itself.");
}
if (from.rootAgent().findAgent(name).isEmpty()) {
return new IllegalArgumentException(
"Agent '"
+ from.name()
+ "' requested a transfer to an agent that is not in the agent tree.");
}
return new IllegalArgumentException(
"Agent '" + from.name() + "' is not allowed to transfer to the requested agent.");
}

/**
* Runs {@code callEvent}'s function calls and emits what follows from them: the tool-confirmation
* request, the function response, and any structured final response. Shared by the normal
Expand Down
106 changes: 106 additions & 0 deletions core/src/test/java/com/google/adk/agents/LlmAgentTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -835,6 +835,112 @@ public void runAsync_resumeAfterTransfer_userEventWithoutFunctionResponse_resume
assertThat(rootLlm.getRequests()).isEmpty();
}

// The transfer belongs to a run of root that already ended, as in a past loop iteration.
@Test
public void runAsync_resumeAfterThisAgentEnded_ignoresTheEarlierTransfer() {
LlmAgent sub =
createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub should not run")))
.name("sub")
.build();
TestLlm rootLlm = createTestLlm(createTextLlmResponse("root runs again"));
LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build();
InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
for (Event event :
ImmutableList.of(
Event.builder()
.id("e1")
.invocationId("inv")
.author("root")
.actions(EventActions.builder().endOfAgent(true).build())
.build(),
Event.builder()
.id("u1")
.invocationId("inv")
.author("user")
.content(Content.fromParts(Part.fromText("next")))
.build())) {
var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
}
context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);
context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true);

List<Event> events = root.runAsync(context).toList().blockingGet();

assertThat(simplifyEvents(events)).contains("root: root runs again");
assertThat(simplifyEvents(events)).doesNotContain("sub: sub should not run");
}

// The target's end marker is from before root's latest transfer, as in a past loop iteration.
@Test
public void runAsync_resumeTransferAfterTargetEndedEarlier_runsTheTarget() {
TestLlm subLlm = createTestLlm(createTextLlmResponse("sub runs"));
LlmAgent sub = createTestAgentBuilder(subLlm).name("sub").build();
LlmAgent root =
createTestAgentBuilder(createTestLlm(createTextLlmResponse("root should not run")))
.name("root")
.subAgents(sub)
.build();
InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
for (Event event :
ImmutableList.of(
Event.builder()
.id("e1")
.invocationId("inv")
.author("sub")
.actions(EventActions.builder().endOfAgent(true).build())
.build(),
Event.builder()
.id("e2")
.invocationId("inv")
.author("root")
.actions(EventActions.builder().endOfAgent(true).build())
.build(),
transferEvent("root", "sub").toBuilder().id("t2").build())) {
var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
}
context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);
context.setAgentState("sub", /* agentState= */ null, /* endOfAgent= */ true);

List<Event> events = root.runAsync(context).toList().blockingGet();

assertThat(simplifyEvents(events)).contains("sub: sub runs");
}

// The deprecated setEndInvocation flag on a content event is not an end marker.
@Test
public void runAsync_resumeAfterContentWithEndFlag_stillResumesTheTransfer() {
LlmAgent sub =
createTestAgentBuilder(createTestLlm(createTextLlmResponse("sub resumes")))
.name("sub")
.build();
TestLlm rootLlm = createTestLlm(createTextLlmResponse("root should not run"));
LlmAgent root = createTestAgentBuilder(rootLlm).name("root").subAgents(sub).build();
InvocationContext context = resumableContextWithSeededEvent(root, transferEvent("root", "sub"));
for (Event event :
ImmutableList.of(
Event.builder()
.id("e1")
.invocationId("inv")
.author("root")
.content(Content.fromParts(Part.fromText("ending")))
.actions(EventActions.builder().endOfAgent(true).build())
.build(),
Event.builder()
.id("u1")
.invocationId("inv")
.author("user")
.content(Content.fromParts(Part.fromText("next")))
.build())) {
var unused = context.sessionService().appendEvent(context.session(), event).blockingGet();
}
context.setAgentState("root", ImmutableMap.of(), /* endOfAgent= */ false);

List<Event> events = root.runAsync(context).toList().blockingGet();

assertThat(simplifyEvents(events)).contains("sub: sub resumes");
assertThat(rootLlm.getRequests()).isEmpty();
}

// Mirrors AgentTransferTest.transferTargets_ordersSubAgentsThenParentThenPeers: the flow keeps a
// second copy of this rule that the agents package cannot reach, so both copies are pinned.
@Test
Expand Down
Loading
Loading