diff --git a/core/src/main/java/com/google/adk/workflow/ResumeScan.java b/core/src/main/java/com/google/adk/workflow/ResumeScan.java new file mode 100644 index 000000000..431f0e285 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ResumeScan.java @@ -0,0 +1,509 @@ +/* + * 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.workflow; + +import static com.google.common.base.Strings.isNullOrEmpty; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectReader; +import com.google.adk.JsonBaseModel; +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.Role; +import com.google.adk.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Sets; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** + * Rebuilds the progress of workflow nodes from the session events of their invocation, and decides + * how a node that already ran in an earlier turn resumes. {@link #scan} folds the event log into + * {@link RecoveredNode}s; when a workflow or {@link Context#runNode} reaches a node, {@link + * #intercept} decides whether it runs again or replays its earlier result through {@link + * #replayContext}. + */ +@Experimental +final class ResumeScan { + + private static final String RESULT_KEY = "result"; + // Stands for text that is not JSON, since the JSON text null reads as a null value. + private static final Object NOT_JSON = new Object(); + private static final ObjectReader STRICT_JSON = + JsonBaseModel.getMapper() + .readerFor(Object.class) + .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + + private ResumeScan() {} + + /** + * Reconstructs the direct children of the node at {@code basePath} from the session events of the + * invocation of {@code context}, as {@link #scan(List, String, String)} does. + */ + static ImmutableMap scan(InvocationContext context, String basePath) { + return scan(context.session().immutableEvents(), basePath, context.invocationId()); + } + + /** + * Reconstructs the direct children of the node at {@code basePath} from the final events of + * {@code invocationId} in one forward pass, keyed by their {@code name@runId} segment. An empty + * base path reconstructs root nodes. A user event can only answer an interrupt; any other event + * reports the progress of the node that emitted it. + */ + static ImmutableMap scan( + List events, String basePath, String invocationId) { + Reconstruction reconstruction = new Reconstruction(basePath); + for (Event event : events) { + if (!invocationId.equals(event.invocationId()) || event.partial().orElse(false)) { + continue; + } + if (Objects.equals(event.author(), Role.USER)) { + reconstruction.recordAnswers(event); + } else { + reconstruction.recordNodeEvent(event); + } + } + return ImmutableMap.copyOf(reconstruction.recovered); + } + + /** + * Decides whether {@code node} runs again or replays its result from {@code recovered}. In order: + * a node with unanswered interrupts waits on them, unless it reruns on resume and has answers; a + * failed node reruns; a recorded output, route or transfer replays; a node {@code canceledMidRun} + * in this turn reruns; an answered node that does not rerun on resume completes with its answers. + * Otherwise a node that reruns on resume reruns unless it finished since its last answer, and a + * {@code dynamic} node reruns while a graph node does not. + */ + static Interception intercept( + Node node, @Nullable RecoveredNode recovered, boolean dynamic, boolean canceledMidRun) { + if (recovered == null) { + return Interception.run(ImmutableMap.of()); + } + Map resolved = copyOf(recovered.resolvedResponses); + if (node instanceof Workflow) { + // A workflow replays its own children, so it always runs. + return Interception.run(resolved); + } + String transfer = recovered.transferToAgent; + String branch = recovered.branch; + ImmutableSet unresolved = recovered.unresolved(); + + if (!unresolved.isEmpty()) { + return node.rerunOnResume() && !resolved.isEmpty() + ? Interception.run(resolved) + : Interception.waitOn(branch, unresolved, transfer); + } + if (recovered.failed) { + return Interception.run(resolved); + } + Object output = recovered.output; + if (output != null || recovered.route != null || transfer != null) { + return Interception.replay( + output instanceof Content content ? processRehydratedOutput(node, content) : output, + recovered.route, + branch, + transfer); + } + // A node canceled mid-run in this turn left progress events but no result. + if (canceledMidRun) { + return Interception.run(resolved); + } + if (!recovered.interruptIds.isEmpty() && !node.rerunOnResume()) { + Object answers = resolved.size() == 1 ? resolved.values().iterator().next() : resolved; + return Interception.replay(answers, /* route= */ null, branch, /* transferToAgent= */ null); + } + // A rerun node that finished without an output stays done, so its side effects don't repeat. + boolean shouldRun = + node.waitForOutput() || (node.rerunOnResume() ? !recovered.finishedAfterResume : dynamic); + return shouldRun + ? Interception.run(resolved) + : Interception.replay( + /* output= */ null, /* route= */ null, branch, /* transferToAgent= */ null); + } + + /** + * Returns the context that {@code node} would have finished with, carrying the replayed result of + * {@code interception} without running the node. + */ + static Context replayContext( + Context parent, Node node, String runId, boolean useAsOutput, Interception interception) { + InvocationContext invocationContext = + interception.branch() == null + ? parent.invocationContext() + : parent.invocationContext().toBuilder().branch(interception.branch()).build(); + Context context = + new Context( + invocationContext, + NodeExecutionState.builder() + .node(node) + .eventSink(parent.requireNodeState().eventSink()) + .parent(parent) + .runId(runId) + .useAsOutput(useAsOutput) + .build()); + interception.replayOn(context); + return context; + } + + /** + * Returns the output of a node whose content was its output: the content's text, read as JSON + * when the node has an output schema that does not accept a string. Blank text is no output. + * + * @throws IllegalArgumentException if the schema needs JSON and the text is not JSON + */ + static @Nullable Object processRehydratedOutput(Node node, Content content) { + String text = textOf(content).trim(); + if (text.isEmpty()) { + return null; + } + Optional schema = node.outputSchema(); + // Model text is raw rather than JSON-quoted, so it stays text when the schema accepts strings. + if (schema.isEmpty() || acceptsString(schema.get())) { + return text; + } + Object value = readJson(text); + if (value == NOT_JSON) { + throw new IllegalArgumentException( + "Validation failed for rehydrated output against schema: value is not valid JSON."); + } + return value; + } + + /** + * Returns the answer that {@code response} carries: the value of a {@code {"result": value}} + * envelope, read as JSON when it is a string that holds JSON, or else the response itself. + */ + private static @Nullable Object unwrapAnswer(@Nullable Object response) { + if (!(response instanceof Map map) || map.size() != 1 || !map.containsKey(RESULT_KEY)) { + return response; + } + Object value = map.get(RESULT_KEY); + if (!(value instanceof String text)) { + return value; + } + Object json = readJson(text); + return json == NOT_JSON ? text : json; + } + + /** Returns whether {@code schema} accepts a string, directly or through an {@code anyOf}. */ + private static boolean acceptsString(Schema schema) { + return typeOf(schema) == Type.Known.STRING + || schema.anyOf().orElse(ImmutableList.of()).stream().anyMatch(ResumeScan::acceptsString); + } + + /** Reads {@code text} as strict JSON, or returns {@link #NOT_JSON} if it is not JSON. */ + private static @Nullable Object readJson(String text) { + try { + return STRICT_JSON.readValue(text); + } catch (JsonProcessingException e) { + return NOT_JSON; + } + } + + private static Type.@Nullable Known typeOf(Schema schema) { + return schema.type().map(Type::knownEnum).orElse(null); + } + + /** Returns the text of the non-thought parts of {@code content}. */ + private static String textOf(Content content) { + StringBuilder text = new StringBuilder(); + for (Part part : content.parts().orElse(ImmutableList.of())) { + if (!part.thought().orElse(false)) { + part.text().ifPresent(text::append); + } + } + return text.toString(); + } + + private static Map copyOf(Map map) { + return Collections.unmodifiableMap(new LinkedHashMap<>(map)); + } + + /** + * The execution state of a node reconstructed from an earlier turn of its invocation. Only the + * scan that creates it updates it. + */ + static final class RecoveredNode { + @Nullable Object output; + boolean failed; + @Nullable ImmutableList route; + @Nullable String branch; + @Nullable String transferToAgent; + final Set interruptIds = new LinkedHashSet<>(); + final Map resolvedResponses = new LinkedHashMap<>(); + + /** Whether the node's last direct event, since its last answer, shows that it finished. */ + boolean finishedAfterResume; + + /** Returns the IDs of the interrupts the node raised that are not answered yet. */ + ImmutableSet unresolved() { + return Sets.difference(interruptIds, resolvedResponses.keySet()).immutableCopy(); + } + + /** Records {@code answer}; what the node emitted before it asked is not its result. */ + void recordAnswer(String interruptId, @Nullable Object answer) { + output = null; + finishedAfterResume = false; + resolvedResponses.put(interruptId, answer); + } + } + + /** + * How a node that already ran resumes: whether it runs again with {@code resumeInputs}, keyed by + * interrupt ID, or replays its result on {@code branch}, where null keeps its parent's branch. + */ + record Interception( + boolean shouldRun, + @Nullable Object output, + @Nullable ImmutableList route, + @Nullable String branch, + ImmutableSet interrupts, + Map resumeInputs, + @Nullable String transferToAgent) { + + /** The node runs again and reads {@code resumeInputs}. */ + static Interception run(Map resumeInputs) { + return new Interception( + /* shouldRun= */ true, + /* output= */ null, + /* route= */ null, + /* branch= */ null, + ImmutableSet.of(), + resumeInputs, + /* transferToAgent= */ null); + } + + /** The node still waits on its unanswered {@code interrupts}. */ + static Interception waitOn( + @Nullable String branch, + ImmutableSet interrupts, + @Nullable String transferToAgent) { + return new Interception( + /* shouldRun= */ false, + /* output= */ null, + /* route= */ null, + branch, + interrupts, + ImmutableMap.of(), + transferToAgent); + } + + /** The node replays the result it recorded, without running. */ + static Interception replay( + @Nullable Object output, + @Nullable ImmutableList route, + @Nullable String branch, + @Nullable String transferToAgent) { + return new Interception( + /* shouldRun= */ false, + output, + route, + branch, + ImmutableSet.of(), + ImmutableMap.of(), + transferToAgent); + } + + /** Records the replayed result on {@code context} without sending it. */ + void replayOn(Context context) { + NodeExecutionState state = context.requireNodeState(); + if (output != null) { + state.produceOutput(output); + } + if (route != null) { + state.selectRoutes(route); + } + if (transferToAgent != null) { + context.eventActions().setTransferToAgent(transferToAgent); + } + state.addInterruptIds(interrupts); + } + } + + /** + * One pass over an event log. An event belongs to the {@code name@runId} segment of the direct + * child of {@code basePath} that it comes from. + */ + private static final class Reconstruction { + private final String basePath; + private final String childPrefix; + + /** Reconstructed nodes keyed by their owner, in the order they first appear. */ + final Map recovered = new LinkedHashMap<>(); + + /** The node that raised each interrupt, so a later answer finds its owner. */ + private final Map interruptOwner = new HashMap<>(); + + Reconstruction(String basePath) { + this.basePath = basePath; + this.childPrefix = basePath.isEmpty() ? "" : basePath + "/"; + } + + /** Attributes {@code event} to its owner and records what it says about that node. */ + void recordNodeEvent(Event event) { + String eventPath = event.nodeInfo().map(NodeInfo::path).orElse(""); + String key = ownerKey(eventPath); + if (key == null) { + return; + } + String ownerPath = childPrefix + key; + // False for an event that a descendant of the owner emitted. + boolean isDirect = eventPath.equals(ownerPath); + RecoveredNode node = recovered.computeIfAbsent(key, unused -> new RecoveredNode()); + // The finish check below reads the failure and the interrupts that these record. + recordOutcome(event, isDirect, ownerPath, node); + boolean raisedAny = recordRaisedInterrupts(event, node); + if (isDirect) { + node.finishedAfterResume = + !raisedAny + && !node.failed + && event.functionCalls().isEmpty() + && event.functionResponses().isEmpty(); + } + } + + /** + * Records the answers that the function responses of the user event {@code event} carry. A + * response whose ID matches no interrupt can still answer an interrupt raised on a tool + * sub-branch named after that interrupt, such as {@code wf@1.task@ask_1} for {@code ask_1}. + */ + void recordAnswers(Event event) { + for (FunctionResponse response : event.functionResponses()) { + String id = response.id().orElse(null); + if (id == null) { + continue; + } + RecoveredNode owner = interruptOwner.get(id); + ImmutableMap targets = + owner != null + ? ImmutableMap.of(id, owner) + : interruptsMatchingBranch(event.branch().orElse(null)); + if (targets.isEmpty()) { + continue; + } + Object answer = unwrapAnswer(response.response().orElse(null)); + targets.forEach((interruptId, node) -> node.recordAnswer(interruptId, answer)); + } + } + + /** Returns the raised interrupts whose IDs are run IDs of {@code branch}'s segments. */ + private ImmutableMap interruptsMatchingBranch(@Nullable String branch) { + if (isNullOrEmpty(branch)) { + return ImmutableMap.of(); + } + Set branchRunIds = new HashSet<>(); + for (String segment : Splitter.on('.').split(branch)) { + int at = segment.lastIndexOf('@'); + if (at >= 0 && at < segment.length() - 1) { + branchRunIds.add(segment.substring(at + 1)); + } + } + Map targets = new LinkedHashMap<>(); + for (RecoveredNode node : recovered.values()) { + for (String interruptId : node.interruptIds) { + if (branchRunIds.contains(interruptId)) { + targets.put(interruptId, node); + } + } + } + return ImmutableMap.copyOf(targets); + } + + /** Returns the key that owns an event at {@code eventPath}, or null if it is outside. */ + private @Nullable String ownerKey(String eventPath) { + if (eventPath.equals(basePath) || !eventPath.startsWith(childPrefix)) { + return null; + } + String rest = eventPath.substring(childPrefix.length()); + int slash = rest.indexOf('/'); + return slash < 0 ? rest : rest.substring(0, slash); + } + + /** + * Records the output, routes, transfer and failure that {@code event} shows for the node at + * {@code ownerPath}. A descendant's event counts only if the descendant's output is also the + * owner's, as with {@code useAsOutput}. + */ + private static void recordOutcome( + Event event, boolean isDirect, String ownerPath, RecoveredNode node) { + Object eventOutput = event.output().orElse(null); + if (eventOutput == null && event.nodeInfo().map(NodeInfo::messageAsOutput).orElse(false)) { + eventOutput = event.content().orElse(null); + } + boolean isDelegated = + eventOutput != null + && event.nodeInfo().map(info -> info.outputFor().contains(ownerPath)).orElse(false); + if (!isDirect && !isDelegated) { + return; + } + if (eventOutput != null) { + node.output = eventOutput; + if (event.output().isPresent()) { + node.branch = event.branch().orElse(null); + } + } + Optional> route = event.actions().route(); + Optional transfer = event.actions().transferToAgent(); + route.ifPresent(routes -> node.route = routes); + transfer.ifPresent(agent -> node.transferToAgent = agent); + // A result or a later event without an error clears a failure; a failure drops the result. + boolean hasResult = eventOutput != null || route.isPresent() || transfer.isPresent(); + if (hasResult || event.errorCode().isEmpty()) { + node.failed = false; + } else { + node.output = null; + node.route = null; + node.transferToAgent = null; + node.failed = true; + } + } + + /** + * Records the long-running call interrupts that {@code event} raises for {@code node}, and + * returns whether {@code event} raised any. + */ + private boolean recordRaisedInterrupts(Event event, RecoveredNode node) { + Set raised = event.longRunningToolIds().orElse(ImmutableSet.of()); + for (String id : raised) { + // An interrupt raised again waits for a new answer, as in Python. + node.resolvedResponses.remove(id); + node.interruptIds.add(id); + interruptOwner.put(id, node); + } + return !raised.isEmpty(); + } + } +} diff --git a/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java new file mode 100644 index 000000000..b64115532 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java @@ -0,0 +1,358 @@ +/* + * 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.workflow; + +import static com.google.adk.testing.WorkflowTestNodes.dispatcher; +import static com.google.adk.testing.WorkflowTestNodes.emitter; +import static com.google.adk.workflow.WorkflowResumeTesting.AGE_SCHEMA; +import static com.google.adk.workflow.WorkflowResumeTesting.LONG_RUNNING_CALL; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.callContent; +import static com.google.adk.workflow.WorkflowResumeTesting.modelText; +import static com.google.adk.workflow.WorkflowResumeTesting.newEvent; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.events.Event; +import com.google.adk.workflow.ResumeScan.Interception; +import com.google.adk.workflow.ResumeScan.RecoveredNode; +import com.google.adk.workflow.WorkflowResumeTesting.MessageOutputNode; +import com.google.adk.workflow.WorkflowResumeTesting.WaitingNode; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Flowable; +import java.util.Map; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ResumeScan}. */ +@RunWith(JUnit4.class) +public final class ResumeScanTest { + + /** An object schema with an integer {@code count}. */ + private static final Schema COUNT_SCHEMA = + Schema.builder() + .type(Type.Known.OBJECT) + .properties(ImmutableMap.of("count", Schema.builder().type(Type.Known.INTEGER).build())) + .build(); + + @Test + public void scan_interruptRaisedBelowAChild_isAnsweredOnThatChild() { + ImmutableList events = + ImmutableList.of( + interruptEvent("wf@1/d@1/c@c", "ask_1"), answerEvent("ask_1", ImmutableMap.of("v", 1))); + + ImmutableMap recovered = ResumeScan.scan(events, "wf@1", "inv"); + + assertThat(recovered.keySet()).containsExactly("d@1"); + assertThat(recovered.get("d@1").resolvedResponses) + .containsExactly("ask_1", ImmutableMap.of("v", 1)); + } + + @Test + public void scan_answerFromAnotherInvocation_resolvesNothing() { + ImmutableList events = + ImmutableList.of( + interruptEvent("wf@1/ask@1", "ask_1"), + answerEvent("ask_1", ImmutableMap.of("v", 1)).toBuilder() + .invocationId("other") + .build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("ask@1"); + + assertThat(recovered.unresolved()).containsExactly("ask_1"); + } + + @Test + public void scan_interruptRaisedAgainAfterItsAnswer_waitsForANewAnswer() { + ImmutableList events = + ImmutableList.of( + interruptEvent("wf@1/ask@1", "ask_1"), + answerEvent("ask_1", ImmutableMap.of("v", 1)), + interruptEvent("wf@1/ask@1", "ask_1")); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("ask@1"); + + assertThat(recovered.unresolved()).containsExactly("ask_1"); + } + + @Test + public void scan_outputOnAnEventWithAnErrorCode_replaysTheOutput() { + ImmutableList events = + ImmutableList.of( + nodeEvent("wf@1/chat@1") + .output("cut short") + .errorCode(new FinishReason("MAX_TOKENS")) + .build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("chat@1"); + Interception interception = + ResumeScan.intercept( + emitter("chat", "fresh"), recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + + assertThat(recovered.failed).isFalse(); + assertThat(interception.shouldRun()).isFalse(); + assertThat(interception.output()).isEqualTo("cut short"); + } + + @Test + public void scan_partialEvent_isIgnored() { + ImmutableList events = + ImmutableList.of( + nodeEvent("wf@1/chat@1").errorCode(new FinishReason("RuntimeError")).build(), + nodeEvent("wf@1/chat@1") + .partial(true) + .content(Content.builder().role("model").parts(Part.fromText("partial")).build()) + .nodeInfo(NodeInfo.builder().path("wf@1/chat@1").messageAsOutput(true).build()) + .build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("chat@1"); + + assertThat(recovered.output).isNull(); + assertThat(recovered.failed).isTrue(); + assertThat(recovered.finishedAfterResume).isFalse(); + } + + @Test + public void scan_answerOwnedByAnotherInterrupt_doesNotResolveByTheBranchsRunId() { + ImmutableList events = + ImmutableList.of( + interruptEvent("wf@1/a@1", "1"), + interruptEvent("wf@1/b@1", "ask_b"), + answerEvent("ask_b", ImmutableMap.of("result", "ok")).toBuilder() + .branch("wf@1.b@1") + .build()); + + ImmutableMap recovered = ResumeScan.scan(events, "wf@1", "inv"); + + assertThat(recovered.get("a@1").resolvedResponses).isEmpty(); + assertThat(recovered.get("b@1").resolvedResponses).containsExactly("ask_b", "ok"); + } + + @Test + public void scan_answerOnABranchNamedAfterAnInterrupt_resolvesThatInterrupt() { + ImmutableList events = + ImmutableList.of( + interruptEvent("wf@1/ask@1", "ask_1"), + answerEvent("call_9", ImmutableMap.of("result", "ok")).toBuilder() + .branch("wf@1.task@ask_1") + .build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("ask@1"); + + assertThat(recovered.resolvedResponses).containsExactly("ask_1", "ok"); + } + + @Test + public void scan_failureAfterAnOutput_dropsTheOutput() { + ImmutableList events = + ImmutableList.of( + nodeEvent("wf@1/n@1").output("stale").build(), + nodeEvent("wf@1/n@1").errorCode(new FinishReason("IllegalStateException")).build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("n@1"); + + assertThat(recovered.output).isNull(); + assertThat(recovered.failed).isTrue(); + } + + @Test + public void intercept_finishedRerunNodeThatWaitsForOutput_stillRuns() { + Node node = new WaitingNode("n", (ctx, input) -> Flowable.empty(), /* rerunOnResume= */ true); + RecoveredNode answered = new RecoveredNode(); + answered.interruptIds.add("a"); + answered.resolvedResponses.put("a", 1); + answered.finishedAfterResume = true; + RecoveredNode unasked = new RecoveredNode(); + unasked.finishedAfterResume = true; + + Interception answeredInterception = + ResumeScan.intercept(node, answered, /* dynamic= */ false, /* canceledMidRun= */ false); + Interception unaskedInterception = + ResumeScan.intercept(node, unasked, /* dynamic= */ false, /* canceledMidRun= */ false); + + assertThat(answeredInterception.shouldRun()).isTrue(); + assertThat(unaskedInterception.shouldRun()).isTrue(); + } + + @Test + public void intercept_nothingRecorded_runsADynamicNodeButSkipsAGraphNode() { + Node node = emitter("n", "out"); + RecoveredNode recovered = new RecoveredNode(); + + Interception dynamic = + ResumeScan.intercept(node, recovered, /* dynamic= */ true, /* canceledMidRun= */ false); + Interception graph = + ResumeScan.intercept(node, recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + + assertThat(dynamic.shouldRun()).isTrue(); + assertThat(graph.shouldRun()).isFalse(); + } + + @Test + public void intercept_nodeCanceledMidRun_runsAgainAsAGraphNode() { + RecoveredNode progressOnly = new RecoveredNode(); + + Interception interception = + ResumeScan.intercept( + emitter("n", "out"), progressOnly, /* dynamic= */ false, /* canceledMidRun= */ true); + + assertThat(interception.shouldRun()).isTrue(); + } + + @Test + public void intercept_answeredNodeThatDoesNotRerun_completesWithAllAnswersAsAMap() { + RecoveredNode recovered = new RecoveredNode(); + recovered.interruptIds.addAll(ImmutableList.of("a", "b")); + recovered.resolvedResponses.put("a", 1); + recovered.resolvedResponses.put("b", 2); + + Interception interception = + ResumeScan.intercept( + emitter("n", "out"), recovered, /* dynamic= */ false, /* canceledMidRun= */ false); + + assertThat(interception.shouldRun()).isFalse(); + assertThat(interception.output()).isEqualTo(ImmutableMap.of("a", 1, "b", 2)); + } + + @Test + public void intercept_unresolvedInterruptWithPartialProgress_runsOnlyANodeThatRerunsOnResume() { + RecoveredNode partialProgress = new RecoveredNode(); + partialProgress.interruptIds.addAll(ImmutableList.of("a", "b")); + partialProgress.resolvedResponses.put("a", 1); + RecoveredNode noProgress = new RecoveredNode(); + noProgress.interruptIds.addAll(ImmutableList.of("a", "b")); + Node rerunNode = dispatcher("rerun", (ctx, input) -> Flowable.empty()); + Node graphNode = emitter("static", "out"); + + Interception rerunWithProgress = + ResumeScan.intercept( + rerunNode, partialProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + Interception rerunWithoutProgress = + ResumeScan.intercept( + rerunNode, noProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + Interception staticWithProgress = + ResumeScan.intercept( + graphNode, partialProgress, /* dynamic= */ false, /* canceledMidRun= */ false); + + assertThat(rerunWithProgress.shouldRun()).isTrue(); + assertThat(rerunWithProgress.resumeInputs()).containsExactly("a", 1); + assertThat(rerunWithoutProgress.shouldRun()).isFalse(); + assertThat(rerunWithoutProgress.interrupts()).containsExactly("a", "b"); + assertThat(staticWithProgress.shouldRun()).isFalse(); + assertThat(staticWithProgress.interrupts()).containsExactly("b"); + } + + @Test + public void processRehydratedOutput_outputNoLongerFitsTheSchema_resumesUnvalidated() { + Schema schema = + Schema.builder() + .type(Type.Known.OBJECT) + .properties(ImmutableMap.of("city", Schema.builder().type(Type.Known.STRING).build())) + .required("city") + .build(); + + Object output = + ResumeScan.processRehydratedOutput( + new MessageOutputNode("chat", "", schema), modelText("{\"town\": \"Paris\"}")); + + assertThat(output).isEqualTo(ImmutableMap.of("town", "Paris")); + } + + @Test + public void processRehydratedOutput_textThatIsNotJsonForAnObjectSchema_fails() { + Node node = new MessageOutputNode("chat", "", AGE_SCHEMA); + Content notJsonText = modelText("not json"); + Content bareWordText = modelText("hello"); + + IllegalArgumentException notJson = + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.processRehydratedOutput(node, notJsonText)); + IllegalArgumentException bareWord = + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.processRehydratedOutput(node, bareWordText)); + + assertThat(notJson) + .hasMessageThat() + .contains("Validation failed for rehydrated output against schema"); + assertThat(bareWord) + .hasMessageThat() + .contains("Validation failed for rehydrated output against schema"); + } + + @Test + public void processRehydratedOutput_blankText_isNoOutput() { + Object output = + ResumeScan.processRehydratedOutput( + new MessageOutputNode("chat", "", COUNT_SCHEMA), modelText(" ")); + + assertThat(output).isNull(); + } + + @Test + public void processRehydratedOutput_jsonNullForAnObjectSchema_isNoOutput() { + Object output = + ResumeScan.processRehydratedOutput( + new MessageOutputNode("chat", "", AGE_SCHEMA), modelText("null")); + + assertThat(output).isNull(); + } + + @Test + public void processRehydratedOutput_stringSchema_keepsTheText() { + Node node = new MessageOutputNode("chat", "", Schema.builder().type(Type.Known.STRING).build()); + + assertThat(ResumeScan.processRehydratedOutput(node, modelText("42"))).isEqualTo("42"); + } + + /** + * Returns an event in which the node at {@code path} waits on the long-running call {@code + * interruptId}. + */ + private static Event interruptEvent(String path, String interruptId) { + return nodeEvent(path) + .content(callContent(LONG_RUNNING_CALL, interruptId)) + .longRunningToolIds(ImmutableSet.of(interruptId)) + .build(); + } + + /** A user's answer event as a session holds it. */ + private static Event answerEvent(String interruptId, Map response) { + return Event.builder() + .id(Event.generateEventId()) + .author("user") + .invocationId("inv") + .content(answer(interruptId, response)) + .build(); + } + + private static Event.Builder nodeEvent(String path) { + return newEvent() + .author("wf") + .invocationId("inv") + .nodeInfo(NodeInfo.builder().path(path).build()); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java new file mode 100644 index 000000000..234b49ed8 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/WorkflowResumeTesting.java @@ -0,0 +1,118 @@ +/* + * 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.workflow; + +import com.google.adk.agents.Context; +import com.google.adk.events.Event; +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 com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Flowable; +import java.util.Map; +import java.util.Optional; +import java.util.function.BiFunction; +import org.jspecify.annotations.Nullable; + +/** Helpers for the tests that pause a workflow and resume it. */ +final class WorkflowResumeTesting { + // The name of the long-running call that the nodes of these tests wait on. + static final String LONG_RUNNING_CALL = "approve"; + + /** An object schema that requires an integer {@code age}. */ + static final Schema AGE_SCHEMA = + Schema.builder() + .type(Type.Known.OBJECT) + .properties(ImmutableMap.of("age", Schema.builder().type(Type.Known.INTEGER).build())) + .required("age") + .build(); + + private WorkflowResumeTesting() {} + + static Content answer(String interruptId, Map response) { + return answer(interruptId, response, LONG_RUNNING_CALL); + } + + @SuppressWarnings("unchecked") // A function response's map holds JSON values. + static Content answer(String interruptId, Map response, String name) { + return Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id(interruptId) + .name(name) + .response((Map) response) + .build()) + .build()) + .build(); + } + + static Content callContent(String name, String id) { + return Content.builder() + .role("model") + .parts( + Part.builder().functionCall(FunctionCall.builder().name(name).id(id).build()).build()) + .build(); + } + + static Content modelText(String text) { + return Content.builder().role("model").parts(Part.fromText(text)).build(); + } + + static Event.Builder newEvent() { + return Event.builder().id(Event.generateEventId()).author(""); + } + + /** A node that waits for an output and runs its {@code body}. */ + record WaitingNode( + String name, BiFunction> body, boolean rerunOnResume) + implements Node { + + @Override + public boolean waitForOutput() { + return true; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return body.apply(context, nodeInput); + } + } + + /** A node whose model-style message is its output, as an agent node's can be. */ + record MessageOutputNode(String name, String text, @Nullable Schema schema) implements Node { + + @Override + public Optional outputSchema() { + return Optional.ofNullable(schema); + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Flowable.just( + newEvent() + .content(modelText(text)) + .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) + .build()); + } + } +}