Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
11 changes: 10 additions & 1 deletion core/src/main/java/com/google/adk/agents/ParallelAgent.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,14 @@

import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
import com.google.adk.events.Event;
import com.google.common.collect.ImmutableSet;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import io.reactivex.rxjava3.core.Flowable;
import io.reactivex.rxjava3.core.Scheduler;
import io.reactivex.rxjava3.schedulers.Schedulers;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

Expand Down Expand Up @@ -175,13 +177,16 @@ protected Flowable<Event> runAsyncImpl(InvocationContext invocationContext) {
return Flowable.empty();
}

ImmutableSet<String> directSubAgentNames =
currentSubAgents.stream().map(BaseAgent::name).collect(ImmutableSet.toImmutableSet());

var updatedInvocationContext = setBranchForCurrentAgent(this, invocationContext);
List<Flowable<Event>> agentFlowables = new ArrayList<>();
for (BaseAgent subAgent : currentSubAgents) {
agentFlowables.add(subAgent.runAsync(updatedInvocationContext).subscribeOn(scheduler));
}
return Flowable.merge(agentFlowables)
.takeUntil((Event event) -> event.actions().escalate().orElse(false));
.takeUntil((Event event) -> asksThisAgentToExit(event, directSubAgentNames));
}

/**
Expand All @@ -195,4 +200,8 @@ protected Flowable<Event> runLiveImpl(InvocationContext invocationContext) {
return Flowable.error(
new UnsupportedOperationException("runLive is not defined for ParallelAgent yet."));
}

private static boolean asksThisAgentToExit(Event event, Set<String> directSubAgentNames) {
return event.actions().escalate().orElse(false) && directSubAgentNames.contains(event.author());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -134,4 +134,56 @@ public void runAsync_escalationEvent_shortCircuitsOtherAgents() {
// Test RxJava Disposal behavior: SlowAgent won't emit anything
subscriber.assertValueCount(2);
}

@Test
public void runAsync_nestedLoopEscalation_keepsSiblingBranchesRunning() {
TestScheduler testScheduler = new TestScheduler();

TestAgent escalatingAgent =
new TestAgent(
"escalating_agent",
10,
testScheduler,
"Escalating!",
EventActions.builder().escalate(true).build());

TestAgent slowAgent = new TestAgent("slow_agent", 100, testScheduler, "Finished");

LoopAgent loopAgent =
LoopAgent.builder().name("loop").subAgents(escalatingAgent).maxIterations(3).build();

ParallelAgent parallelAgent =
ParallelAgent.builder()
.name("parallel_agent")
.subAgents(loopAgent, slowAgent)
.scheduler(testScheduler)
.build();

InvocationContext invocationContext = createInvocationContext(parallelAgent);

var subscriber = parallelAgent.runAsync(invocationContext).test();

// Escalation is raised on the first iteration, so the loop stops there even though
// maxIterations(3) would have allowed two more passes at 20ms and 30ms. Advancing
// past all three windows proves the cut came from the escalation, not the cap.
testScheduler.advanceTimeBy(40, MILLISECONDS);
subscriber.assertValueCount(1);
assertThat(subscriber.values().get(0).author()).isEqualTo("escalating_agent");
// The escalation came from a nested agent, not a direct sub-agent, so the parallel
// agent must not short-circuit its remaining branches.
subscriber.assertNotComplete();

// Slow agent completes at 100ms
testScheduler.advanceTimeBy(100, MILLISECONDS);
subscriber.assertValueCount(2);

Event event1 = subscriber.values().get(0);
assertThat(event1.author()).isEqualTo("escalating_agent");
assertThat(event1.actions().escalate()).hasValue(true);

Event event2 = subscriber.values().get(1);
assertThat(event2.author()).isEqualTo("slow_agent");

subscriber.assertComplete();
}
}