From 22bdfa760000bee9a4b8b761d1910a50805d4771 Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Sun, 12 Apr 2026 03:08:14 -0700 Subject: [PATCH] fix: the RemoteA2AAgent does not support resumability and always transfers to the parent Only LLMAgent supported this via the disallowTransferToParent attribute Add Resumable interface and test for runner resumability. This change introduces the `Resumable` interface, allowing agents to indicate if they can be resumed from a previous state. Also the LLMAgent and RemoteA2A agent has been modified to implement this. Ea A new test case in `RunnerResumabilityTest` demonstrates how the runner can leverage this to resume execution from a resumable sub-agent based on the session history. PiperOrigin-RevId: 898502514 --- .../google/adk/a2a/agent/RemoteA2AAgent.java | 23 ++- .../adk/a2a/agent/RemoteA2AAgentTest.java | 12 ++ .../java/com/google/adk/agents/LlmAgent.java | 7 +- .../java/com/google/adk/agents/Resumable.java | 21 +++ .../java/com/google/adk/runner/Runner.java | 9 +- .../adk/runner/RunnerResumabilityTest.java | 141 ++++++++++++++++++ .../com/google/adk/testing/TestUtils.java | 4 + 7 files changed, 203 insertions(+), 14 deletions(-) create mode 100644 core/src/main/java/com/google/adk/agents/Resumable.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java diff --git a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java index d4e094710..6a76a28e8 100644 --- a/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java +++ b/a2a/src/main/java/com/google/adk/a2a/agent/RemoteA2AAgent.java @@ -27,6 +27,7 @@ import com.google.adk.agents.BaseAgent; import com.google.adk.agents.Callbacks; import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.Resumable; import com.google.adk.events.Event; import com.google.adk.utils.AgentEnums.AgentOrigin; import com.google.common.collect.ImmutableList; @@ -75,7 +76,7 @@ *
  • Converting A2A client responses back into ADK format * */ -public class RemoteA2AAgent extends BaseAgent { +public class RemoteA2AAgent extends BaseAgent implements Resumable { private static final Logger logger = LoggerFactory.getLogger(RemoteA2AAgent.class); private static final ObjectMapper objectMapper = @@ -85,6 +86,7 @@ public class RemoteA2AAgent extends BaseAgent { private final Client a2aClient; private String description; private final boolean streaming; + private final boolean resumable; // Internal constructor used by builder private RemoteA2AAgent(Builder builder) { @@ -118,6 +120,7 @@ private RemoteA2AAgent(Builder builder) { this.description = this.agentCard.description(); } this.streaming = builder.streaming && this.agentCard.capabilities().streaming(); + this.resumable = builder.resumable; } public static Builder builder() { @@ -134,6 +137,7 @@ public static class Builder { private List beforeAgentCallback; private List afterAgentCallback; private boolean streaming; + private boolean resumable = true; @CanIgnoreReturnValue public Builder streaming(boolean streaming) { @@ -141,6 +145,12 @@ public Builder streaming(boolean streaming) { return this; } + @CanIgnoreReturnValue + public Builder resumable(boolean resumable) { + this.resumable = resumable; + return this; + } + @CanIgnoreReturnValue public Builder name(String name) { this.name = name; @@ -192,6 +202,11 @@ public boolean isStreaming() { return streaming; } + @Override + public boolean isResumable() { + return resumable; + } + private Message.Builder newA2AMessage(Message.Role role, List> parts) { return new Message.Builder().messageId(UUID.randomUUID().toString()).role(role).parts(parts); } @@ -227,8 +242,7 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { return Flowable.create( emitter -> { StreamHandler handler = - new StreamHandler( - emitter.serialize(), invocationContext, requestJson, streaming, name()); + new StreamHandler(emitter.serialize(), invocationContext, requestJson, name()); ImmutableList> consumers = ImmutableList.of(handler::handleEvent); a2aClient.sendMessage(originalMessage, consumers, handler::handleError, null); @@ -249,7 +263,6 @@ private static class StreamHandler { private final FlowableEmitter emitter; private final InvocationContext invocationContext; private final String requestJson; - private final boolean streaming; private final String agentName; private boolean done = false; private final StringBuilder textBuffer = new StringBuilder(); @@ -259,12 +272,10 @@ private static class StreamHandler { FlowableEmitter emitter, InvocationContext invocationContext, String requestJson, - boolean streaming, String agentName) { this.emitter = emitter; this.invocationContext = invocationContext; this.requestJson = requestJson; - this.streaming = streaming; this.agentName = agentName; } diff --git a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java index 0609c3b04..495ed33ea 100644 --- a/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java +++ b/a2a/src/test/java/com/google/adk/a2a/agent/RemoteA2AAgentTest.java @@ -127,6 +127,18 @@ public void createAgent_streaming_true_returnsStreamingAgent() { assertThat(agent.isStreaming()).isTrue(); } + @Test + public void createAgent_resumable_default_true() { + RemoteA2AAgent agent = getAgentBuilder().build(); + assertThat(agent.isResumable()).isTrue(); + } + + @Test + public void createAgent_resumable_false() { + RemoteA2AAgent agent = getAgentBuilder().resumable(false).build(); + assertThat(agent.isResumable()).isFalse(); + } + @Test public void runAsync_aggregatesPartialEvents() { RemoteA2AAgent agent = createAgent(); 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 98bba4606..e6e2c6734 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -76,7 +76,7 @@ import org.slf4j.LoggerFactory; /** The LLM-based agent. */ -public class LlmAgent extends BaseAgent { +public class LlmAgent extends BaseAgent implements Resumable { private static final Logger logger = LoggerFactory.getLogger(LlmAgent.class); @@ -779,6 +779,11 @@ public boolean disallowTransferToParent() { return disallowTransferToParent; } + @Override + public boolean isResumable() { + return !disallowTransferToParent(); + } + public boolean disallowTransferToPeers() { return disallowTransferToPeers; } diff --git a/core/src/main/java/com/google/adk/agents/Resumable.java b/core/src/main/java/com/google/adk/agents/Resumable.java new file mode 100644 index 000000000..d2967ef36 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/Resumable.java @@ -0,0 +1,21 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.agents; + +/** Interface for agents that can be resumed from history directly. */ +public interface Resumable { + boolean isResumable(); +} diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 44a281f72..729fbb5d0 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -22,6 +22,7 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.Resumable; import com.google.adk.agents.RunConfig; import com.google.adk.apps.App; import com.google.adk.artifacts.BaseArtifactService; @@ -752,13 +753,7 @@ protected Flowable runLiveImpl( private boolean isTransferableAcrossAgentTree(BaseAgent agentToRun) { BaseAgent current = agentToRun; while (current != null) { - // Agents eligible to transfer must have an LLM-based agent parent. - if (!(current instanceof LlmAgent)) { - return false; - } - // If any agent can't transfer to its parent, the chain is broken. - LlmAgent agent = (LlmAgent) current; - if (agent.disallowTransferToParent()) { + if (!(current instanceof Resumable resumableAgent) || !resumableAgent.isResumable()) { return false; } current = current.parentAgent(); diff --git a/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java new file mode 100644 index 000000000..701e11ca9 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerResumabilityTest.java @@ -0,0 +1,141 @@ +package com.google.adk.runner; + +import static com.google.adk.testing.TestUtils.createContent; +import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.Resumable; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.sessions.Session; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Flowable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests for the resumability feature in the {@link Runner}. + * + *

    This test suite verifies that when a session contains events from a resumable agent, the + * runner correctly bypasses earlier workflow steps (like the root agent) and resumes execution + * directly at the last active resumable agent. + */ +@RunWith(JUnit4.class) +public final class RunnerResumabilityTest { + + private static class TestResumableAgent extends BaseAgent implements Resumable { + private final boolean resumable; + + public TestResumableAgent(String name, boolean resumable) { + super(name, "", ImmutableList.of(), ImmutableList.of(), ImmutableList.of()); + this.resumable = resumable; + } + + @Override + public boolean isResumable() { + return resumable; + } + + @Override + protected Flowable runAsyncImpl(InvocationContext context) { + return Flowable.just( + Event.builder() + .id("event-" + name()) + .author(name()) + .content( + Content.builder() + .parts( + ImmutableList.of(Part.builder().text("response from " + name()).build())) + .build()) + .build()); + } + + @Override + protected Flowable runLiveImpl(InvocationContext context) { + return runAsyncImpl(context); + } + } + + /** + * Verifies that {@link Runner#runAsync} picks up execution from a resumable sub-agent. + * + *

    This test configures a workflow hierarchy with a root agent and a resumable sub-agent. By + * pre-loading the session state with an initial event authored by the sub-agent, we simulate a + * scenario where execution stopped within the sub-workflow. Calling {@code runAsync} triggers the + * resume behavior, bypassing the default root flow. + */ + @Test + public void runAsync_resumesAtResumableSubAgent() { + TestResumableAgent subAgent = new TestResumableAgent("sub_agent", true); + LlmAgent rootAgent = + LlmAgent.builder() + .name("root_agent") + .model(createTestLlm(createLlmResponse(createContent("from root")))) + .subAgents(ImmutableList.of(subAgent)) + .build(); + + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(rootAgent).build()).build(); + + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Event subAgentEvent = + Event.builder() + .id("initial-event") + .author("sub_agent") + .content(createContent("subagent greeting")) + .build(); + + var unused = runner.sessionService().appendEvent(session, subAgentEvent).blockingGet(); + + var events = + runner.runAsync("user", session.id(), createContent("continue")).toList().blockingGet(); + + assertThat(simplifyEvents(events)).containsExactly("sub_agent: response from sub_agent"); + } + + /** + * Verifies that {@link Runner#runAsync} does not resume execution from a non-resumable sub-agent. + * + *

    This test ensures that if the sub-agent is NOT marked as resumable, even if there are + * existing events from it in the session, the runner falls back to regular execution starting + * from the root agent. + */ + @Test + public void runAsync_doesNotResumeAtNonResumableSubAgent() { + TestResumableAgent subAgent = new TestResumableAgent("sub_agent", false); + LlmAgent rootAgent = + LlmAgent.builder() + .name("root_agent") + .model(createTestLlm(createLlmResponse(createContent("from root")))) + .subAgents(ImmutableList.of(subAgent)) + .build(); + + Runner runner = + Runner.builder().app(App.builder().name("test").rootAgent(rootAgent).build()).build(); + + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + Event subAgentEvent = + Event.builder() + .id("initial-event") + .author("sub_agent") + .content(createContent("subagent greeting")) + .build(); + + var unused = runner.sessionService().appendEvent(session, subAgentEvent).blockingGet(); + + var events = + runner.runAsync("user", session.id(), createContent("continue")).toList().blockingGet(); + + assertThat(simplifyEvents(events)).containsExactly("root_agent: from root"); + } +} diff --git a/core/src/test/java/com/google/adk/testing/TestUtils.java b/core/src/test/java/com/google/adk/testing/TestUtils.java index daed8d2e4..2f56be5dc 100644 --- a/core/src/test/java/com/google/adk/testing/TestUtils.java +++ b/core/src/test/java/com/google/adk/testing/TestUtils.java @@ -234,6 +234,10 @@ public static TestLlm createTestLlm(Supplier> responsesSup return new TestLlm(responsesSupplier); } + public static Content createContent(String text) { + return Content.builder().parts(Part.builder().text(text).build()).build(); + } + public static LlmResponse createLlmResponse(Content content) { return LlmResponse.builder().content(content).build(); }