From f5c75618992b8a3cfbc6aedf84bc4e0e7ddf11ba Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Mon, 21 Sep 2026 01:15:54 -0700 Subject: [PATCH] feat: honor the plain-text continuation flag and allow only one resumability flag `ResumabilityConfig.plainTextContinuationAutoResume` previously had no effect. It now selects the same resumption behavior as `resumable`, and building a config with both flags set throws `IllegalArgumentException`. PiperOrigin-RevId: 985095038 --- .../google/adk/agents/InvocationContext.java | 8 +- .../google/adk/apps/ResumabilityConfig.java | 30 +- .../java/com/google/adk/runner/Runner.java | 10 +- .../adk/agents/InvocationContextTest.java | 39 ++ .../adk/apps/ResumabilityConfigTest.java | 65 +++ .../runner/RunnerLegacyResumabilityTest.java | 417 ++++++++++++++++++ .../com/google/adk/runner/RunnerTest.java | 327 +------------- 7 files changed, 571 insertions(+), 325 deletions(-) create mode 100644 core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java create mode 100644 core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java 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 26a8800c6..384fd8372 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -306,10 +306,14 @@ public Optional contextCacheConfig() { /** * Returns whether the current invocation is resumable. Mirrors Python ADK v1's {@code - * InvocationContext.is_resumable}. + * InvocationContext.is_resumable}. The deprecated plain-text continuation shim selects the same + * resumption behavior, so it reports resumable too; the two are mutually exclusive. */ + @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. public boolean isResumable() { - return resumabilityConfig != null && resumabilityConfig.isResumable(); + return resumabilityConfig != null + && (resumabilityConfig.isResumable() + || resumabilityConfig.isPlainTextContinuationAutoResume()); } private static class InvocationCostManager { diff --git a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java index d6d0445d7..590d9177a 100644 --- a/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java +++ b/core/src/main/java/com/google/adk/apps/ResumabilityConfig.java @@ -16,6 +16,8 @@ package com.google.adk.apps; +import static com.google.common.base.Preconditions.checkArgument; + import com.google.auto.value.AutoValue; import com.google.errorprone.annotations.CanIgnoreReturnValue; @@ -23,6 +25,9 @@ * App resumability config, mirroring Python ADK v1's {@code ResumabilityConfig}: pause on a * long-running call and resume from the last event. Applies to all agents in the app. * + *

The two flags select the same resumption behavior and are mutually exclusive: set {@link + * #isResumable()}, or the deprecated shim, but not both. + * * @deprecated Partial feature: only event-reconstruction-based pause/resume for {@code * SequentialAgent} is implemented. Full session resumability (persisted agent state, durable * resume, other workflow agents) is not yet available. Forward-compatible: the same config will @@ -41,10 +46,13 @@ public abstract class ResumabilityConfig { * default, matching Python ADK, where a plain-text {@code runAsync} always starts a new * invocation and a paused invocation is resumed explicitly. * + *

Selects the same resumption behavior as {@link #isResumable()}, with which it is mutually + * exclusive; it differs only in also resuming on a plain-text continuation. + * * @deprecated Back-compat shim for callers that deliver a resume as a plain-text turn. Migrate to * {@code Runner.runAsync(userId, sessionId, invocationId, message, runConfig, stateDelta)} - * (or send a function response to the paused call) and stop setting this flag; it will be - * removed. + * (or send a function response to the paused call) and set {@link #isResumable()} instead; + * this flag will be removed. */ @Deprecated public abstract boolean isPlainTextContinuationAutoResume(); @@ -70,6 +78,22 @@ public abstract static class Builder { @CanIgnoreReturnValue public abstract Builder plainTextContinuationAutoResume(boolean value); - public abstract ResumabilityConfig build(); + abstract ResumabilityConfig autoBuild(); + + /** + * Builds the config, rejecting a combination of flags that has no defined behavior. + * + * @throws IllegalArgumentException if both resumability and the deprecated shim are set; they + * select the same behavior, so exactly one may be enabled. + */ + @SuppressWarnings("deprecation") // Validating the deprecated shim against the supported flag. + public ResumabilityConfig build() { + ResumabilityConfig config = autoBuild(); + checkArgument( + !(config.isResumable() && config.isPlainTextContinuationAutoResume()), + "resumable and plainTextContinuationAutoResume are mutually exclusive: set resumable for" + + " the supported flag, or the deprecated shim, but not both."); + return config; + } } } 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 fbecbf02b..2642033a3 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -863,9 +863,15 @@ private boolean isTransferableAcrossAgentTree(BaseAgent agentToRun) { return true; } - /** Returns whether resumability is enabled for this runner's app. */ + /** + * Returns whether resumability is enabled for this runner's app, by either the supported flag or + * the deprecated plain-text continuation shim, which selects the same behavior. + */ + @SuppressWarnings("deprecation") // The shim it reads is deprecated by design. private boolean isResumable() { - return resumabilityConfig != null && resumabilityConfig.isResumable(); + return resumabilityConfig != null + && (resumabilityConfig.isResumable() + || resumabilityConfig.isPlainTextContinuationAutoResume()); } /** Returns the agent that should handle the next request based on session history. */ 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 55f9b4d65..995387157 100644 --- a/core/src/test/java/com/google/adk/agents/InvocationContextTest.java +++ b/core/src/test/java/com/google/adk/agents/InvocationContextTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertThrows; import static org.mockito.Mockito.mock; +import com.google.adk.apps.ResumabilityConfig; import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.events.Event; import com.google.adk.memory.BaseMemoryService; @@ -73,6 +74,44 @@ public void setUp() { activeStreamingTools.put("test-tool", new ActiveStreamingTool(new LiveRequestQueue())); } + // The deprecated shim selects the same resumption behavior as resumable(true), so every + // resumability branch keyed on isResumable() must treat it identically. + @Test + public void isResumable_shimOnly_reportsResumable() { + InvocationContext shimContext = contextWith(resumabilityConfigWithShim()); + InvocationContext resumableContext = contextWith(resumabilityConfigResumable()); + InvocationContext neitherContext = contextWith(null); + + assertThat(shimContext.isResumable()).isTrue(); + assertThat(resumableContext.isResumable()).isTrue(); + assertThat(neitherContext.isResumable()).isFalse(); + } + + @SuppressWarnings("deprecation") // Exercises the deprecated shim. + private static ResumabilityConfig resumabilityConfigWithShim() { + return ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); + } + + @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. + private static ResumabilityConfig resumabilityConfigResumable() { + return ResumabilityConfig.builder().resumable(true).build(); + } + + @SuppressWarnings("deprecation") // ResumabilityConfig is deprecated until durable resumability. + private InvocationContext contextWith(ResumabilityConfig resumabilityConfig) { + return InvocationContext.builder() + .sessionService(mockSessionService) + .artifactService(mockArtifactService) + .pluginManager(pluginManager) + .invocationId(testInvocationId) + .agent(mockAgent) + .session(session) + .userContent(userContent) + .runConfig(runConfig) + .resumabilityConfig(resumabilityConfig) + .build(); + } + @Test public void testBuildWithUserContent() { InvocationContext context = diff --git a/core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java b/core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java new file mode 100644 index 000000000..2407f5c62 --- /dev/null +++ b/core/src/test/java/com/google/adk/apps/ResumabilityConfigTest.java @@ -0,0 +1,65 @@ +/* + * 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.apps; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ResumabilityConfig}. */ +@RunWith(JUnit4.class) +@SuppressWarnings("deprecation") // Exercises the deprecated plainTextContinuationAutoResume shim. +public final class ResumabilityConfigTest { + + @Test + public void build_defaults_selectNoResumption() { + ResumabilityConfig config = ResumabilityConfig.builder().build(); + + assertThat(config.isResumable()).isFalse(); + assertThat(config.isPlainTextContinuationAutoResume()).isFalse(); + } + + @Test + public void build_resumableOnly_succeeds() { + ResumabilityConfig config = ResumabilityConfig.builder().resumable(true).build(); + + assertThat(config.isResumable()).isTrue(); + assertThat(config.isPlainTextContinuationAutoResume()).isFalse(); + } + + @Test + public void build_legacyShimOnly_succeeds() { + ResumabilityConfig config = + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build(); + + assertThat(config.isResumable()).isFalse(); + assertThat(config.isPlainTextContinuationAutoResume()).isTrue(); + } + + @Test + public void build_bothFlags_throws() { + ResumabilityConfig.Builder builder = + ResumabilityConfig.builder().resumable(true).plainTextContinuationAutoResume(true); + + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, builder::build); + + assertThat(thrown).hasMessageThat().contains("mutually exclusive"); + } +} diff --git a/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java new file mode 100644 index 000000000..7f0302758 --- /dev/null +++ b/core/src/test/java/com/google/adk/runner/RunnerLegacyResumabilityTest.java @@ -0,0 +1,417 @@ +/* + * Copyright 2025 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.runner; + +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.TestUtils.simplifyEvents; +import static com.google.common.truth.Truth.assertThat; + +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.LoopAgent; +import com.google.adk.agents.ParallelAgent; +import com.google.adk.agents.SequentialAgent; +import com.google.adk.apps.App; +import com.google.adk.apps.ResumabilityConfig; +import com.google.adk.events.Event; +import com.google.adk.flows.llmflows.Functions; +import com.google.adk.sessions.Session; +import com.google.adk.telemetry.Tracing; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.FunctionTool; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +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 io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; +import io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import org.jspecify.annotations.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Runner tests for the deprecated plain-text continuation shim. + * + *

Includes a copy of every resumability test that predates durable checkpoints, re-run under the + * shim, so the shim keeps behaving exactly as resumability did before it was split in two. + */ +@RunWith(JUnit4.class) +public final class RunnerLegacyResumabilityTest { + @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); + + private Tracer originalTracer; + + @Before + public void setUp() { + this.originalTracer = Tracing.getTracer(); + Tracing.setTracerForTesting( + openTelemetryRule.getOpenTelemetry().getTracer("RunnerLegacyResumabilityTest")); + } + + @After + public void tearDown() { + Tracing.setTracerForTesting(originalTracer); + } + + public static class Tools { + private Tools() {} + + public static ImmutableMap echoTool(String message) { + return ImmutableMap.of("message", message); + } + + // A long-running tool awaiting an external result has nothing to return yet; FunctionTool + // coerces the absent return into an empty response. + @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. + public static @Nullable ImmutableMap pendingTool(String message) { + return null; + } + + static final AtomicInteger pendingProgressToolCalls = new AtomicInteger(0); + + // A long-running tool that reports progress: it returns a non-empty "pending" status on the + // initial call. Counts executions so a test can assert it runs exactly once across turns. + @SuppressWarnings("unused") // Invoked reflectively by FunctionTool. + public static ImmutableMap pendingProgressTool(String message) { + pendingProgressToolCalls.incrementAndGet(); + return ImmutableMap.of("status", "pending"); + } + } + + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void + runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // With resumability on, B pauses right after requesting confirmation (no extra model call), so + // a + // single follow-up response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("Response after user confirmed.")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools(FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ true)) + .build(); + LlmAgent agentC = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) + .name("c_agent") + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(workflowAgent) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List eventsBeforeConfirmation = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Turn 1: A runs, B pauses for confirmation, and C must not run yet. + assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); + + FunctionCall askUserConfirmationFunctionCall = + Iterables.getOnlyElement( + eventsBeforeConfirmation.stream() + .map(Functions::getAskUserConfirmationFunctionCalls) + .filter(functionCalls -> !functionCalls.isEmpty()) + .findFirst() + .get()); + List eventsAfterConfirmation = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(askUserConfirmationFunctionCall.id().get()) + .name(askUserConfirmationFunctionCall.name().get()) + .response(ImmutableMap.of("confirmed", true))) + .build())) + .toList() + .blockingGet(); + + // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. + assertThat(simplifyEvents(eventsAfterConfirmation)) + .containsExactly( + "b_agent: FunctionResponse(name=echoTool, response={message=hello})", + "b_agent: Response after user confirmed.", + "c_agent: agent C done") + .inOrder(); + } + + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void + runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume_legacyShim() { + LlmAgent agentA = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) + .name("a_agent") + .build(); + // With resumability on, B pauses right after the long-running call (no extra model call), so a + // single follow-up response covers the resume. + TestLlm bTestLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("agent B resumed")); + LlmAgent agentB = + createTestAgentBuilder(bTestLlm) + .name("b_agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent agentC = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) + .name("c_agent") + .build(); + SequentialAgent workflowAgent = + SequentialAgent.builder() + .name("workflow_agent") + .subAgents(ImmutableList.of(agentA, agentB, agentC)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(workflowAgent) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List eventsBeforeResume = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. B must not + // make + // a further model call after the pending call. + assertThat(simplifyEvents(eventsBeforeResume)).contains("a_agent: agent A done"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); + assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); + + List eventsAfterResume = + runner + .runAsync( + "user", + session.id(), + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("lro_call_id") + .name("echoTool") + .response(ImmutableMap.of("message", "hello"))) + .build())) + .toList() + .blockingGet(); + + // Turn 2: B resumes from the long-running response, then C runs. A is not re-run. + assertThat(simplifyEvents(eventsAfterResume)) + .containsExactly("b_agent: agent B resumed", "c_agent: agent C done") + .inOrder(); + } + + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall_legacyShim() { + TestLlm testLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + // Extra responses the flow must NOT consume; reaching them means it looped. + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("should not be reached")); + LlmAgent agent = + createTestAgentBuilder(testLlm) + .name("agent") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(agent) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // The flow paused after the single long-running call instead of re-calling the model. + assertThat(testLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); + } + + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void + runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration_legacyShim() { + AtomicInteger calls = new AtomicInteger(); + TestLlm loopLlm = + createTestLlm( + () -> + calls.incrementAndGet() <= 5 + ? Flowable.just( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello"))) + : Flowable.just(createTextLlmResponse("stop"))); + LlmAgent inner = + createTestAgentBuilder(loopLlm) + .name("inner") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LoopAgent loop = + LoopAgent.builder() + .name("loop") + .subAgents(ImmutableList.of(inner)) + .maxIterations(3) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(loop) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List unused = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // Paused after the first iteration: one model call, not maxIterations. + assertThat(loopLlm.getRequests()).hasSize(1); + } + + @Test + @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). + public void + runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes_legacyShim() { + TestLlm longRunningLlm = + createTestLlm( + createFunctionCallLlmResponse( + "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), + createTextLlmResponse("unexpected")); + LlmAgent longRunningBranch = + createTestAgentBuilder(longRunningLlm) + .name("long_running_branch") + .tools( + FunctionTool.create( + Tools.class, + "echoTool", + /* requireConfirmation= */ false, + /* isLongRunning= */ true)) + .build(); + LlmAgent plainBranch = + createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) + .name("plain_branch") + .build(); + ParallelAgent parallel = + ParallelAgent.builder() + .name("parallel") + .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) + .build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name("test") + .rootAgent(parallel) + .resumabilityConfig( + ResumabilityConfig.builder().plainTextContinuationAutoResume(true).build()) + .build()) + .build(); + Session session = runner.sessionService().createSession("test", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) + .toList() + .blockingGet(); + + // The long-running branch paused after one model call; the other branch still completed. + assertThat(longRunningLlm.getRequests()).hasSize(1); + assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); + } +} diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 3870d3461..b7adda6dd 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -45,8 +45,6 @@ import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; -import com.google.adk.agents.LoopAgent; -import com.google.adk.agents.ParallelAgent; import com.google.adk.agents.RunConfig; import com.google.adk.agents.SequentialAgent; import com.google.adk.apps.App; @@ -122,14 +120,21 @@ public final class RunnerTest { @Rule public final OpenTelemetryRule openTelemetryRule = OpenTelemetryRule.create(); private final BasePlugin plugin = mockPlugin("test"); + private final Content pluginContent = createContent("from plugin"); + private final TestLlm testLlm = createTestLlm(createLlmResponse(createContent("from llm"))); + private final LlmAgent agent = createTestAgentBuilder(testLlm).build(); + private Runner runner; + private Session session; + private Tracer originalTracer; private final FailingEchoTool failingEchoTool = new FailingEchoTool(); + private final EchoTool echoTool = new EchoTool(); private final TestLlm testLlmWithFunctionCall = @@ -2353,219 +2358,6 @@ public void runAsync_withToolConfirmation_inSequentialAgentSubAgent_resumesSubAg .inOrder(); } - // OSS HITL: after an adk_request_confirmation resumes sub-agent B in a SequentialAgent(A, B, C), - // the workflow must advance to C without re-running the already completed A. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withToolConfirmation_inSequentialAgent_runsLaterSubAgentsAfterResume() { - LlmAgent agentA = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) - .name("a_agent") - .build(); - // With resumability on, B pauses right after requesting confirmation (no extra model call), so - // a - // single follow-up response covers the resume. - TestLlm bTestLlm = - createTestLlm( - createFunctionCallLlmResponse( - "tool_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("Response after user confirmed.")); - LlmAgent agentB = - createTestAgentBuilder(bTestLlm) - .name("b_agent") - .tools(FunctionTool.create(Tools.class, "echoTool", /* requireConfirmation= */ true)) - .build(); - LlmAgent agentC = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) - .name("c_agent") - .build(); - SequentialAgent workflowAgent = - SequentialAgent.builder() - .name("workflow_agent") - .subAgents(ImmutableList.of(agentA, agentB, agentC)) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List eventsBeforeConfirmation = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // Turn 1: A runs, B pauses for confirmation, and C must not run yet. - assertThat(simplifyEvents(eventsBeforeConfirmation)).contains("a_agent: agent A done"); - assertThat(simplifyEvents(eventsBeforeConfirmation)).doesNotContain("c_agent: agent C done"); - - FunctionCall askUserConfirmationFunctionCall = - Iterables.getOnlyElement( - eventsBeforeConfirmation.stream() - .map(Functions::getAskUserConfirmationFunctionCalls) - .filter(functionCalls -> !functionCalls.isEmpty()) - .findFirst() - .get()); - List eventsAfterConfirmation = - runner - .runAsync( - "user", - session.id(), - Content.fromParts( - Part.builder() - .functionResponse( - FunctionResponse.builder() - .id(askUserConfirmationFunctionCall.id().get()) - .name(askUserConfirmationFunctionCall.name().get()) - .response(ImmutableMap.of("confirmed", true))) - .build())) - .toList() - .blockingGet(); - - // Turn 2: B resumes and executes the tool, then C runs. A is not re-run. - assertThat(simplifyEvents(eventsAfterConfirmation)) - .containsExactly( - "b_agent: FunctionResponse(name=echoTool, response={message=hello})", - "b_agent: Response after user confirmed.", - "c_agent: agent C done") - .inOrder(); - } - - // Long-running-call HITL: a pending long-running function call (not the confirmation flow) pauses - // SequentialAgent(A, B, C) after B; on resume B continues and C runs, without re-running A. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withLongRunningCall_inSequentialAgent_runsLaterSubAgentsAfterResume() { - LlmAgent agentA = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent A done"))) - .name("a_agent") - .build(); - // With resumability on, B pauses right after the long-running call (no extra model call), so a - // single follow-up response covers the resume. - TestLlm bTestLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("agent B resumed")); - LlmAgent agentB = - createTestAgentBuilder(bTestLlm) - .name("b_agent") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - LlmAgent agentC = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("agent C done"))) - .name("c_agent") - .build(); - SequentialAgent workflowAgent = - SequentialAgent.builder() - .name("workflow_agent") - .subAgents(ImmutableList.of(agentA, agentB, agentC)) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(workflowAgent) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List eventsBeforeResume = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // Turn 1: A runs, B issues the long-running call and pauses; C must not run yet. B must not - // make - // a further model call after the pending call. - assertThat(simplifyEvents(eventsBeforeResume)).contains("a_agent: agent A done"); - assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("b_agent: agent B resumed"); - assertThat(simplifyEvents(eventsBeforeResume)).doesNotContain("c_agent: agent C done"); - - List eventsAfterResume = - runner - .runAsync( - "user", - session.id(), - Content.fromParts( - Part.builder() - .functionResponse( - FunctionResponse.builder() - .id("lro_call_id") - .name("echoTool") - .response(ImmutableMap.of("message", "hello"))) - .build())) - .toList() - .blockingGet(); - - // Turn 2: B resumes from the long-running response, then C runs. A is not re-run. - assertThat(simplifyEvents(eventsAfterResume)) - .containsExactly("b_agent: agent B resumed", "c_agent: agent C done") - .inOrder(); - } - - // Regression: a pending long-running call must pause the LLM flow after a single model call when - // resumability is on. Before the flow-level pause, the flow kept re-calling the model (re-issuing - // the call), burning tokens. The scripted model would re-issue the call if the flow did not - // pause; - // we assert exactly one model call was made and the later responses were never consumed. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_withLongRunningCall_resumable_pausesAfterSingleModelCall() { - TestLlm testLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - // Extra responses the flow must NOT consume; reaching them means it looped. - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("should not be reached")); - LlmAgent agent = - createTestAgentBuilder(testLlm) - .name("agent") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(agent) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // The flow paused after the single long-running call instead of re-calling the model. - assertThat(testLlm.getRequests()).hasSize(1); - assertThat(simplifyEvents(events)).doesNotContain("agent: should not be reached"); - } - // Gating: with resumability OFF (default) the flow does NOT pause on a long-running call; it // keeps // calling the model as before. Pairs with the resumable test above. @@ -2754,109 +2546,6 @@ private static List resumeWithStatus(Runner runner, Session session, Stri .blockingGet(); } - // A pending long-running call must stop a resumable LoopAgent after the current iteration rather - // than looping again (re-calling the model every iteration), matching Python ADK v1. - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_loopAgentWithLongRunningSubAgent_resumable_stopsAfterFirstIteration() { - AtomicInteger calls = new AtomicInteger(); - TestLlm loopLlm = - createTestLlm( - () -> - calls.incrementAndGet() <= 5 - ? Flowable.just( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello"))) - : Flowable.just(createTextLlmResponse("stop"))); - LlmAgent inner = - createTestAgentBuilder(loopLlm) - .name("inner") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - LoopAgent loop = - LoopAgent.builder() - .name("loop") - .subAgents(ImmutableList.of(inner)) - .maxIterations(3) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(loop) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List unused = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // Paused after the first iteration: one model call, not maxIterations. - assertThat(loopLlm.getRequests()).hasSize(1); - } - - // In a resumable ParallelAgent, a pending long-running call pauses only its own branch (via the - // flow); other branches still complete. ParallelAgent needs no special handling, matching Python - // ADK v1 (cancelling siblings would diverge). - @Test - @SuppressWarnings("deprecation") // Resumability flag is intentionally deprecated (partial). - public void runAsync_parallelAgentWithLongRunningBranch_resumable_otherBranchCompletes() { - TestLlm longRunningLlm = - createTestLlm( - createFunctionCallLlmResponse( - "lro_call_id", "echoTool", ImmutableMap.of("message", "hello")), - createTextLlmResponse("unexpected")); - LlmAgent longRunningBranch = - createTestAgentBuilder(longRunningLlm) - .name("long_running_branch") - .tools( - FunctionTool.create( - Tools.class, - "echoTool", - /* requireConfirmation= */ false, - /* isLongRunning= */ true)) - .build(); - LlmAgent plainBranch = - createTestAgentBuilder(createTestLlm(createTextLlmResponse("plain branch done"))) - .name("plain_branch") - .build(); - ParallelAgent parallel = - ParallelAgent.builder() - .name("parallel") - .subAgents(ImmutableList.of(longRunningBranch, plainBranch)) - .build(); - Runner runner = - Runner.builder() - .app( - App.builder() - .name("test") - .rootAgent(parallel) - .resumabilityConfig(ResumabilityConfig.builder().resumable(true).build()) - .build()) - .build(); - Session session = runner.sessionService().createSession("test", "user").blockingGet(); - - List events = - runner - .runAsync("user", session.id(), Content.fromParts(Part.fromText("from user"))) - .toList() - .blockingGet(); - - // The long-running branch paused after one model call; the other branch still completed. - assertThat(longRunningLlm.getRequests()).hasSize(1); - assertThat(simplifyEvents(events)).contains("plain_branch: plain branch done"); - } - // Resumability disabled (default): a SequentialAgent(A, B, C) does not pause on B's HITL call, so // all sub-agents run in the same turn — matching Python ADK v1 with resumability disabled. @Test @@ -3073,7 +2762,9 @@ public void runner_executesSaveArtifactFlow() { } private static final String BLOB_MIME_TYPE = "example/octet-stream"; + private static final String BLOB_PAYLOAD = "blob payload"; + private static final String PLACEHOLDER_FORMAT = "Uploaded file: %s. It has been saved to the artifacts";