diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 456758b95..481053bd9 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -17,9 +17,12 @@ package com.google.adk.agents; import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.ImmutableSet.toImmutableSet; import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; +import com.google.adk.events.Event; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.plugins.Plugin; @@ -27,8 +30,13 @@ import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import java.util.Collections; import java.util.Map; import java.util.Objects; import java.util.Optional; @@ -41,6 +49,8 @@ @SuppressWarnings("deprecation") // Plumbs the deprecated ResumabilityConfig. public class InvocationContext { + private static final String USER_AUTHOR = "user"; + private final BaseSessionService sessionService; private final BaseArtifactService artifactService; private final BaseMemoryService memoryService; @@ -156,6 +166,80 @@ public Session session() { return session; } + /** + * Returns a snapshot of the session's events, keeping only those on the branch this invocation is + * running on. + * + *

The rule is author-asymmetric on purpose, so a confirmation the user answered on a + * sub-branch stays visible while a descendant agent's own events do not. + */ + public ImmutableList eventsOnCurrentBranch() { + ImmutableList events; + synchronized (session.events()) { + events = ImmutableList.copyOf(session.events()); + } + // Snapshot the mutable branch too, so it cannot change between the id set and the filter. + @Nullable String scopeBranch = branch; + // Only the user-response cross-check needs these, and a null or empty branch skips it. + ImmutableSet branchFunctionCallIds = + isNullOrEmpty(scopeBranch) ? ImmutableSet.of() : branchFunctionCallIds(events, scopeBranch); + return events.stream() + .filter(event -> isOnCurrentBranch(event, scopeBranch, branchFunctionCallIds)) + .collect(toImmutableList()); + } + + /** + * Returns whether {@code event} belongs to this invocation's branch. + * + *

A user event matches this branch, a descendant sub-branch, or no branch at all; one carrying + * function responses must additionally answer a call issued on this branch or below, which is + * what stops a reply leaking in from a parallel tree. Any other event must sit on exactly this + * branch, so a descendant's own events stay hidden. + */ + private boolean isOnCurrentBranch( + Event event, @Nullable String scopeBranch, ImmutableSet branchFunctionCallIds) { + @Nullable String eventBranch = event.branch().orElse(null); + if (!Objects.equals(event.author(), USER_AUTHOR)) { + return Objects.equals(eventBranch, scopeBranch); + } + if (!isNullOrEmpty(scopeBranch)) { + ImmutableSet responseIds = + event.functionResponses().stream() + .map(FunctionResponse::id) + .flatMap(Optional::stream) + .collect(toImmutableSet()); + if (!responseIds.isEmpty() && Collections.disjoint(responseIds, branchFunctionCallIds)) { + return false; + } + } + // Mirrors Python's `self.branch` guard: an empty branch has no descendants. + return eventBranch == null + || scopeBranch == null + || eventBranch.equals(scopeBranch) + || (!scopeBranch.isEmpty() && eventBranch.startsWith(scopeBranch + ".")); + } + + /** + * Returns the IDs of function calls issued on this branch or on a descendant sub-branch. + * + *

Branches are dot-joined, so the trailing dot keeps the prefix test on a segment boundary. + */ + private ImmutableSet branchFunctionCallIds( + ImmutableList events, String scopeBranch) { + String descendantPrefix = scopeBranch + "."; + return events.stream() + .filter( + event -> { + @Nullable String eventBranch = event.branch().orElse(null); + return !isNullOrEmpty(eventBranch) + && (eventBranch.equals(scopeBranch) || eventBranch.startsWith(descendantPrefix)); + }) + .flatMap(event -> event.functionCalls().stream()) + .map(FunctionCall::id) + .flatMap(Optional::stream) + .collect(toImmutableSet()); + } + /** Returns the user content that triggered this invocation, if any. */ public Optional userContent() { return Optional.ofNullable(userContent); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.java b/core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.java index 06e50e76d..d3199bfc8 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessor.java @@ -61,7 +61,8 @@ public class RequestConfirmationLlmRequestProcessor implements RequestProcessor @Override public Single processRequest( InvocationContext invocationContext, LlmRequest llmRequest) { - ImmutableList events = ImmutableList.copyOf(invocationContext.session().events()); + // A confirmation is answered on the branch that asked for it; a parallel tree's is not ours. + ImmutableList events = invocationContext.eventsOnCurrentBranch(); if (events.isEmpty()) { logger.trace( "No events are present in the session. Skipping request confirmation processing."); diff --git a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java index e588a38ca..55f9b4d65 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -21,17 +21,23 @@ import static org.mockito.Mockito.mock; import com.google.adk.artifacts.BaseArtifactService; +import com.google.adk.events.Event; import com.google.adk.memory.BaseMemoryService; import com.google.adk.models.LlmCallsLimitExceededException; import com.google.adk.plugins.PluginManager; import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.Session; import com.google.adk.summarizer.EventsCompactionConfig; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; import java.util.HashMap; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.jspecify.annotations.Nullable; import org.junit.Assert; import org.junit.Before; import org.junit.Test; @@ -724,4 +730,184 @@ public void build_missingSessionService_throwsException() { IllegalStateException exception = assertThrows(IllegalStateException.class, builder::build); assertThat(exception).hasMessageThat().isEqualTo("Session service must be set."); } + + @Test + public void eventsOnCurrentBranch_userEventOnSubBranch_isIncluded() { + Event userOnChild = userEvent("agent_1.child"); + + assertThat(contextOnBranch("agent_1", userOnChild).eventsOnCurrentBranch()) + .containsExactly(userOnChild); + } + + @Test + public void eventsOnCurrentBranch_agentEventOnSubBranch_isExcluded() { + // Asymmetric with the user case on purpose: descendants' internal events stay hidden. + Event agentOnChild = agentEvent("agent_1.child"); + + assertThat(contextOnBranch("agent_1", agentOnChild).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_siblingBranch_isExcluded() { + Event userOnSibling = userEvent("agent_2"); + + assertThat(contextOnBranch("agent_1", userOnSibling).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_userEventOnLookalikeBranch_isExcluded() { + // "agent_10" shares a prefix with "agent_1" but is not a sub-branch of it. + Event userOnLookalike = userEvent("agent_10"); + + assertThat(contextOnBranch("agent_1", userOnLookalike).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_emptyBranch_doesNotMatchBranchedEvents() { + // An empty string is a real branch value, not a synonym for "match everything". + Event userOnBranch = userEvent("agent_1"); + + assertThat(contextOnBranch("", userOnBranch).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_noBranch_matchesEveryUserEventButNotAgentEvents() { + Event userElsewhere = userEvent("agent_2.child"); + Event agentElsewhere = agentEvent("agent_2.child"); + + assertThat(contextOnBranch(null, userElsewhere, agentElsewhere).eventsOnCurrentBranch()) + .containsExactly(userElsewhere); + } + + @Test + public void eventsOnCurrentBranch_userResponseToCallInSubtree_isKept() { + Event callOnChild = callEvent("agent_1.child", "fc_1"); + Event reply = userResponseEvent("agent_1", "fc_1"); + + assertThat(contextOnBranch("agent_1", callOnChild, reply).eventsOnCurrentBranch()) + .containsExactly(reply); + } + + @Test + public void eventsOnCurrentBranch_userResponseToCallElsewhere_isDropped() { + // Sitting on this branch is not enough: the reply answers a parallel tree's call. + Event callElsewhere = callEvent("agent_2", "fc_1"); + Event reply = userResponseEvent("agent_1", "fc_1"); + + assertThat(contextOnBranch("agent_1", callElsewhere, reply).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_userResponseToLookalikeBranchCall_isDropped() { + // "agent_10" shares a prefix with "agent_1" but is not a sub-branch of it. + Event callOnLookalike = callEvent("agent_10", "fc_1"); + Event reply = userResponseEvent("agent_1", "fc_1"); + + assertThat(contextOnBranch("agent_1", callOnLookalike, reply).eventsOnCurrentBranch()) + .isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_severalReplies_eachJudgedAgainstItsOwnCall() { + Event callHere = callEvent("agent_1", "fc_here"); + Event callOnChild = callEvent("agent_1.child", "fc_child"); + Event callElsewhere = callEvent("agent_2", "fc_far"); + Event replyHere = userResponseEvent("agent_1", "fc_here"); + Event replyFar = userResponseEvent("agent_1", "fc_far"); + Event replyChild = userResponseEvent("agent_1", "fc_child"); + + InvocationContext context = + contextOnBranch( + "agent_1", callHere, callOnChild, callElsewhere, replyHere, replyFar, replyChild); + + // callHere matches exactly so it survives; the sub-branch calls do not, but their replies do. + assertThat(context.eventsOnCurrentBranch()) + .containsExactly(callHere, replyHere, replyChild) + .inOrder(); + } + + @Test + public void eventsOnCurrentBranch_emptyBranchAndDotPrefixedEvent_isExcluded() { + // Pins the empty-branch guard: without it the prefix test would admit a dot-prefixed branch. + Event dotPrefixed = userEvent(".x"); + + assertThat(contextOnBranch("", dotPrefixed).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_rootAgentEventWhileOnSubBranch_isExcluded() { + // The narrowing direction: the old scan admitted every event, this one demands equality. + Event rootAgentEvent = agentEvent(null); + + assertThat(contextOnBranch("agent_1", rootAgentEvent).eventsOnCurrentBranch()).isEmpty(); + } + + @Test + public void eventsOnCurrentBranch_rootUserEventWhileOnSubBranch_isIncluded() { + // The user twin of the case above: a null-branch user event still matches. + Event rootUserEvent = userEvent(null); + + assertThat(contextOnBranch("agent_1", rootUserEvent).eventsOnCurrentBranch()) + .containsExactly(rootUserEvent); + } + + @Test + public void eventsOnCurrentBranch_userResponseToUnbranchedCall_isDropped() { + // A root-level call contributes no id, so a reply answering only it is dropped. + Event rootCall = callEvent(null, "fc_1"); + Event reply = userResponseEvent("agent_1", "fc_1"); + + assertThat(contextOnBranch("agent_1", rootCall, reply).eventsOnCurrentBranch()).isEmpty(); + } + + private InvocationContext contextOnBranch(@Nullable String branch, Event... events) { + return InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .memoryService(mockMemoryService) + .pluginManager(pluginManager) + .invocationId(testInvocationId) + .branch(branch) + .agent(mockAgent) + .session(Session.builder("test-session-id").events(ImmutableList.copyOf(events)).build()) + .runConfig(runConfig) + .build(); + } + + private static Event userEvent(@Nullable String branch) { + return Event.builder().author("user").branch(branch).build(); + } + + private static Event agentEvent(@Nullable String branch) { + return Event.builder().author("some_agent").branch(branch).build(); + } + + private static Event callEvent(@Nullable String branch, String callId) { + return Event.builder() + .author("some_agent") + .branch(branch) + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id(callId).name("t").build()) + .build())) + .build(); + } + + private static Event userResponseEvent(@Nullable String branch, String callId) { + return Event.builder() + .author("user") + .branch(branch) + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(callId) + .name("t") + .response(ImmutableMap.of()) + .build()) + .build())) + .build(); + } } diff --git a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java index 6bbd9e55b..ea3e524fa 100644 --- a/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/SequentialAgentTest.java @@ -249,6 +249,22 @@ public void resumeSubAgentIndex_noMatchingAuthor_returnsEmpty() { assertThat(WorkflowAgentResumption.resumeSubAgentIndex(context, root.subAgents())).isEmpty(); } + @Test + public void resumeSubAgentIndex_branchedNestedAuthor_returnsThatSubAgentIndex() { + // A Sequential nested inside a Parallel carries a branch a filter would hide. + TestBaseAgent first = createSubAgent("first_agent"); + TestBaseAgent nested = createSubAgent("nested_agent"); + SequentialAgent branch = + SequentialAgent.builder().name("branch_agent").subAgents(ImmutableList.of(nested)).build(); + SequentialAgent root = + SequentialAgent.builder().name("root").subAgents(ImmutableList.of(first, branch)).build(); + + InvocationContext context = + contextResumingBranchedCall(root, "nested_agent", "p.root", "p.root.branch_agent"); + + assertThat(WorkflowAgentResumption.resumeSubAgentIndex(context, root.subAgents())).hasValue(1); + } + // Session ending with a function response that resumes a call authored by callAuthor. private static InvocationContext contextResumingCall(BaseAgent rootAgent, String callAuthor) { InMemorySessionService sessionService = new InMemorySessionService(); @@ -284,4 +300,45 @@ private static InvocationContext contextResumingCall(BaseAgent rootAgent, String var unusedResponse = sessionService.appendEvent(session, responseEvent).blockingGet(); return createInvocationContext(rootAgent, sessionService, session); } + + // As above, but the context sits on contextBranch and the call event on callBranch. + private static InvocationContext contextResumingBranchedCall( + BaseAgent rootAgent, String callAuthor, String contextBranch, String callBranch) { + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession("test_app", "test-user").blockingGet(); + Event callEvent = + Event.builder() + .id("call_event") + .invocationId("invocationId") + .author(callAuthor) + .branch(callBranch) + .content( + Content.fromParts( + Part.builder() + .functionCall(FunctionCall.builder().id("call_id").name("tool").build()) + .build())) + .build(); + Event responseEvent = + Event.builder() + .id("response_event") + .invocationId("invocationId") + .author("user") + .branch(contextBranch) + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call_id") + .name("tool") + .response(ImmutableMap.of()) + .build()) + .build())) + .build(); + var unusedCall = sessionService.appendEvent(session, callEvent).blockingGet(); + var unusedResponse = sessionService.appendEvent(session, responseEvent).blockingGet(); + InvocationContext context = createInvocationContext(rootAgent, sessionService, session); + context.branch(contextBranch); + return context; + } } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessorTest.java b/core/src/test/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessorTest.java index c8da89026..9ca8c767e 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessorTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/RequestConfirmationLlmRequestProcessorTest.java @@ -20,6 +20,7 @@ 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.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import com.google.adk.agents.InvocationContext; @@ -398,6 +399,29 @@ public void testAgentNameMatchesFixtures() { assertThat(createAgentWithEchoTool().name()).isEqualTo(AGENT_NAME); } + @Test + public void runAsync_approvalOnParallelBranch_doesNotCallOriginalFunction() { + // An approval answered in a parallel tree is not this branch's, though it names the call. + LlmAgent agent = createAgentWithEchoTool(); + Session session = sessionWithApprovalOn("agent_1", "agent_2"); + + assertThat(resumedEventsOnBranch(agent, session, "agent_1")).isEmpty(); + } + + @Test + public void runAsync_approvalOnSubBranch_callsOriginalFunction() { + // The user may answer on a descendant sub-branch, so scoping must not break the normal path. + LlmAgent agent = createAgentWithEchoTool(); + Session session = sessionWithApprovalOn("agent_1", "agent_1.child"); + + ImmutableList resumed = resumedEventsOnBranch(agent, session, "agent_1"); + + assertThat(resumed).hasSize(1); + FunctionResponse response = resumed.get(0).functionResponses().get(0); + assertThat(response.id()).hasValue(ORIGINAL_FUNCTION_CALL_ID); + assertThat(response.name()).hasValue(ECHO_TOOL_NAME); + } + private static ImmutableList resumedEvents(LlmAgent agent, Session session) { return ImmutableList.copyOf( processor @@ -432,6 +456,44 @@ private static InvocationContext buildInvocationContext(LlmAgent agent, Session .build(); } + private static InvocationContext buildInvocationContext( + LlmAgent agent, Session session, String branch) { + return InvocationContext.builder() + .pluginManager(new PluginManager()) + .invocationId(InvocationContext.newInvocationContextId()) + .branch(branch) + .agent(agent) + .session(session) + .sessionService(sessionService) + .build(); + } + + /** + * Returns the legitimate lead-up with the agent's events on {@code agentBranch} and the user's + * approval on {@code approvalBranch}. + */ + private static Session sessionWithApprovalOn(String agentBranch, String approvalBranch) { + ImmutableList events = + CONFIRMED_CALL_EVENTS.stream() + .map( + event -> + event.toBuilder() + .branch(event.author().equals("user") ? approvalBranch : agentBranch) + .build()) + .collect(toImmutableList()); + return Session.builder("session_id").events(events).build(); + } + + private static ImmutableList resumedEventsOnBranch( + LlmAgent agent, Session session, String branch) { + return ImmutableList.copyOf( + processor + .processRequest( + buildInvocationContext(agent, session, branch), LlmRequest.builder().build()) + .blockingGet() + .events()); + } + private static LlmAgent createAgentWithEchoTool() { Content contentWithFunctionCall = Content.fromParts(