From 85780688fb0c1cb02f8c89a4fcbd6ec7dc0eba03 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Fri, 23 Oct 2026 06:14:11 -0700 Subject: [PATCH] feat: pause a workflow for user input PiperOrigin-RevId: 1005544178 --- .../java/com/google/adk/agents/Context.java | 6 +- .../com/google/adk/workflow/BaseNode.java | 26 +- .../java/com/google/adk/workflow/Node.java | 7 +- .../adk/workflow/NodeExecutionState.java | 6 +- .../com/google/adk/workflow/RequestInput.java | 127 ++++++++ .../com/google/adk/workflow/ResumeScan.java | 141 ++++++++- .../google/adk/workflow/ResumeScanTest.java | 135 ++++++++ .../google/adk/workflow/WorkflowHitlTest.java | 293 ++++++++++++++++++ 8 files changed, 714 insertions(+), 27 deletions(-) create mode 100644 core/src/main/java/com/google/adk/workflow/RequestInput.java create mode 100644 core/src/test/java/com/google/adk/workflow/WorkflowHitlTest.java diff --git a/core/src/main/java/com/google/adk/agents/Context.java b/core/src/main/java/com/google/adk/agents/Context.java index 37e2cf9d9..b6d1991c3 100644 --- a/core/src/main/java/com/google/adk/agents/Context.java +++ b/core/src/main/java/com/google/adk/agents/Context.java @@ -323,9 +323,9 @@ public final void routes(List routes) { } /** - * Returns the user's answers to the long-running calls that this workflow node or its descendants - * made in an earlier turn, keyed by interrupt ID. A node that reruns on resume reads these - * answers. A null answer has a null value in the map. + * Returns the user's answers to the requests for input and long-running calls that this workflow + * node or its descendants made in an earlier turn, keyed by interrupt ID. A node that reruns on + * resume reads these answers. A null answer has a null value in the map. * * @throws IllegalStateException if this is not the context of a workflow node. */ diff --git a/core/src/main/java/com/google/adk/workflow/BaseNode.java b/core/src/main/java/com/google/adk/workflow/BaseNode.java index 0c3282d32..6438f7239 100644 --- a/core/src/main/java/com/google/adk/workflow/BaseNode.java +++ b/core/src/main/java/com/google/adk/workflow/BaseNode.java @@ -119,21 +119,27 @@ public Optional stateSchema() { /** * Runs the node and emits its events. Each emission of {@link #runNode} that is an {@link Event} - * passes through, and any other value becomes the output of a new event. + * passes through, a {@link RequestInput} becomes the event that asks for it, and any other value + * becomes the output of a new event. */ final Flowable run(Context context, @Nullable Object nodeInput) { InvocationContext invocationContext = checkNotNull(context).invocationContext(); // The node runner fills in the author, as only it knows the node's place in the graph. return Flowable.defer(() -> runNode(context, nodeInput)) .map( - item -> - item instanceof Event event - ? event - : Event.builder() - .id(invocationContext.newUuid()) - .timestamp(invocationContext.now().toEpochMilli()) - .author("") - .output(item) - .build()); + item -> { + if (item instanceof Event event) { + return event; + } + if (item instanceof RequestInput request) { + return request.toEvent(invocationContext); + } + return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) + .author("") + .output(item) + .build(); + }); } } diff --git a/core/src/main/java/com/google/adk/workflow/Node.java b/core/src/main/java/com/google/adk/workflow/Node.java index 2fdb4ef56..b81fdfd47 100644 --- a/core/src/main/java/com/google/adk/workflow/Node.java +++ b/core/src/main/java/com/google/adk/workflow/Node.java @@ -98,9 +98,10 @@ default boolean requiresAllPredecessors() { /** * Runs the node on {@code nodeInput}, the output it receives from its predecessors. * - *

An emitted {@link Event} passes through, and any other emitted value becomes the node's - * output (at most one per run); emit nothing for no output, since RxJava forbids null. The - * wildcard lets an implementation return a {@code Flowable} of its own type. + *

An emitted {@link Event} passes through, an emitted {@link RequestInput} pauses the workflow + * until the user answers it, and any other emitted value becomes the node's output (at most one + * per run); emit nothing for no output, since RxJava forbids null. The wildcard lets an + * implementation return a {@code Flowable} of its own type. */ @Experimental Flowable runNode(Context context, @Nullable Object nodeInput); diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java index 67a0210b5..9c9db650a 100644 --- a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -196,9 +196,9 @@ public ImmutableSet interruptIds() { } /** - * Returns the user's answers to the long-running calls that this node or its descendants raised - * in an earlier turn, keyed by interrupt ID. A null answer has a null value in the map. The map - * is empty unless the node reruns on resume. + * Returns the user's answers to the requests for input and long-running calls that this node or + * its descendants raised in an earlier turn, keyed by interrupt ID. A null answer has a null + * value in the map. The map is empty unless the node reruns on resume. */ public Map resumeInputs() { return resumeInputs; diff --git a/core/src/main/java/com/google/adk/workflow/RequestInput.java b/core/src/main/java/com/google/adk/workflow/RequestInput.java new file mode 100644 index 000000000..8cabf0524 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/RequestInput.java @@ -0,0 +1,127 @@ +/* + * 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.JsonBaseModel; +import com.google.adk.agents.InvocationContext; +import com.google.adk.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.auto.value.AutoValue; +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.Part; +import com.google.genai.types.Schema; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +import org.jspecify.annotations.Nullable; + +/** + * A node's request for user input, which pauses the workflow until the user answers it. A node + * emits it from {@link Node#runNode}, and the runtime records it as a long-running {@value + * #FUNCTION_CALL_NAME} function call whose ID is {@link #interruptId}. The user answers with a + * function response that carries the same ID. + */ +@Experimental +@AutoValue +public abstract class RequestInput { + + /** The name of the function call that carries a request for input. */ + public static final String FUNCTION_CALL_NAME = "adk_request_input"; + + static final String INTERRUPT_ID_KEY = "interruptId"; + static final String MESSAGE_KEY = "message"; + static final String PAYLOAD_KEY = "payload"; + static final String RESPONSE_SCHEMA_KEY = "response_schema"; + + /** Returns the ID of the request, which the answering function response carries. */ + public abstract String interruptId(); + + /** Returns the message to show the user, if any. */ + public abstract Optional message(); + + /** Returns additional data for rendering the request, if any. */ + public abstract Optional payload(); + + /** + * Returns the schema the answer must match, if any. An answer of the form {@code {"result": + * value}} is unwrapped first, and a string value is read as JSON unless the schema accepts a + * string. + */ + public abstract Optional responseSchema(); + + /** Returns a builder whose interrupt ID is a new random UUID. */ + public static Builder builder() { + return new AutoValue_RequestInput.Builder().interruptId(UUID.randomUUID().toString()); + } + + public abstract Builder toBuilder(); + + /** + * Returns the long-running function-call event that records this request. It has the model role, + * so history rewriting keeps it as the call that the answer responds to. + * + * @throws IllegalArgumentException if the payload has no JSON form + */ + Event toEvent(InvocationContext invocationContext) { + Map args = new LinkedHashMap<>(); + args.put(INTERRUPT_ID_KEY, interruptId()); + args.put(MESSAGE_KEY, message().orElse(null)); + // Converted here so a payload with no JSON form fails the asking node, not the session service. + args.put(PAYLOAD_KEY, payload().map(RequestInput::toJsonValue).orElse(null)); + args.put(RESPONSE_SCHEMA_KEY, responseSchema().map(RequestInput::toJsonValue).orElse(null)); + FunctionCall call = + FunctionCall.builder().name(FUNCTION_CALL_NAME).id(interruptId()).args(args).build(); + return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) + .author("") + .content( + Content.builder() + .role("model") + .parts(Part.builder().functionCall(call).build()) + .build()) + .longRunningToolIds(ImmutableSet.of(interruptId())) + .build(); + } + + private static @Nullable Object toJsonValue(Object value) { + return JsonBaseModel.getMapper().convertValue(value, Object.class); + } + + /** Builder for {@link RequestInput}. */ + @AutoValue.Builder + public abstract static class Builder { + + @CanIgnoreReturnValue + public abstract Builder interruptId(String interruptId); + + @CanIgnoreReturnValue + public abstract Builder message(@Nullable String message); + + @CanIgnoreReturnValue + public abstract Builder payload(@Nullable Object payload); + + @CanIgnoreReturnValue + public abstract Builder responseSchema(@Nullable Schema responseSchema); + + public abstract RequestInput build(); + } +} diff --git a/core/src/main/java/com/google/adk/workflow/ResumeScan.java b/core/src/main/java/com/google/adk/workflow/ResumeScan.java index 431f0e285..7d8ca72c5 100644 --- a/core/src/main/java/com/google/adk/workflow/ResumeScan.java +++ b/core/src/main/java/com/google/adk/workflow/ResumeScan.java @@ -33,10 +33,13 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Sets; 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 java.math.BigInteger; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -203,15 +206,32 @@ static Context replayContext( } /** - * 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. + * Unwraps and validates the answer {@code response} to interrupt {@code interruptId}. A {@code + * {"result": value}} envelope is unwrapped, and a string value inside it is read as JSON unless + * {@code schema} accepts a string. + * + * @throws IllegalArgumentException if the answer does not match {@code schema} */ - private static @Nullable Object unwrapAnswer(@Nullable Object response) { + static @Nullable Object validateAnswer( + @Nullable Schema schema, @Nullable Object response, String interruptId) { + Object answer = unwrapAnswer(schema, response); + if (schema == null) { + return answer; + } + String mismatch = valueMismatch(answer, schema); + if (mismatch != null) { + throw new IllegalArgumentException( + "Validation failed for interrupt " + interruptId + ": " + mismatch); + } + return answer; + } + + private static @Nullable Object unwrapAnswer(@Nullable Schema schema, @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)) { + if (!(value instanceof String text) || (schema != null && acceptsString(schema))) { return value; } Object json = readJson(text); @@ -233,6 +253,84 @@ private static boolean acceptsString(Schema schema) { } } + /** + * Returns why {@code value} does not match {@code schema}, or null if it does. The reason names + * parts of the schema but never the value, which comes from the user. + */ + private static @Nullable String valueMismatch(@Nullable Object value, Schema schema) { + if (value == null && schema.nullable().orElse(false)) { + return null; + } + // Check maps before anyOf so each member is validated as an object with its required keys. + if (value instanceof Map map) { + return mapMismatch(map, schema); + } + List anyOf = schema.anyOf().orElse(ImmutableList.of()); + if (!anyOf.isEmpty() + && anyOf.stream().allMatch(member -> valueMismatch(value, member) != null)) { + return "it matches none of the schema's anyOf alternatives"; + } + Type.Known type = typeOf(schema); + if (type == null || type == Type.Known.TYPE_UNSPECIFIED) { + return null; + } + boolean matches = + switch (type) { + case STRING -> value instanceof String; + case INTEGER -> + value instanceof Integer + || value instanceof Long + || value instanceof Short + || value instanceof Byte + || value instanceof BigInteger; + case NUMBER -> value instanceof Number; + case BOOLEAN -> value instanceof Boolean; + case NULL -> value == null; + case ARRAY -> { + if (!(value instanceof Collection items)) { + yield false; + } + Optional itemSchema = schema.items(); + yield itemSchema.isEmpty() + || items.stream().allMatch(item -> valueMismatch(item, itemSchema.get()) == null); + } + case OBJECT -> false; + default -> true; + }; + return matches ? null : "it is not of the schema's type " + type; + } + + /** Returns why {@code map} does not match {@code schema}, or null if it does. */ + private static @Nullable String mapMismatch(Map map, Schema schema) { + Type.Known type = typeOf(schema); + if (type != null && type != Type.Known.OBJECT && type != Type.Known.TYPE_UNSPECIFIED) { + return "the schema's type is " + type + ", not an object"; + } + List anyOf = schema.anyOf().orElse(ImmutableList.of()); + if (!anyOf.isEmpty() && anyOf.stream().allMatch(member -> mapMismatch(map, member) != null)) { + return "it matches none of the schema's anyOf alternatives"; + } + Optional> properties = schema.properties(); + if (properties.isPresent()) { + for (Map.Entry entry : map.entrySet()) { + Schema property = properties.get().get(String.valueOf(entry.getKey())); + // An undeclared key comes from the user, so the reason does not name it. + if (property == null) { + return "it has a key the schema does not declare"; + } + if (valueMismatch(entry.getValue(), property) != null) { + return "property '" + entry.getKey() + "' does not match its schema"; + } + } + } + for (String required : schema.required().orElse(ImmutableList.of())) { + if (!map.containsKey(required)) { + return "required property '" + required + "' is missing"; + } + } + return null; + } + private static Type.@Nullable Known typeOf(Schema schema) { return schema.type().map(Type::knownEnum).orElse(null); } @@ -248,6 +346,15 @@ private static String textOf(Content content) { return text.toString(); } + /** Reads the response schema that an {@value RequestInput#FUNCTION_CALL_NAME} call carried. */ + private static @Nullable Schema readSchema(@Nullable Object value) { + try { + return JsonBaseModel.getMapper().convertValue(value, Schema.class); + } catch (IllegalArgumentException e) { + return null; + } + } + private static Map copyOf(Map map) { return Collections.unmodifiableMap(new LinkedHashMap<>(map)); } @@ -367,6 +474,9 @@ private static final class Reconstruction { /** The node that raised each interrupt, so a later answer finds its owner. */ private final Map interruptOwner = new HashMap<>(); + /** The response schema of each request for input, to validate its answer. */ + private final Map schemas = new HashMap<>(); + Reconstruction(String basePath) { this.basePath = basePath; this.childPrefix = basePath.isEmpty() ? "" : basePath + "/"; @@ -414,7 +524,7 @@ void recordAnswers(Event event) { if (targets.isEmpty()) { continue; } - Object answer = unwrapAnswer(response.response().orElse(null)); + Object answer = validateAnswer(schemas.get(id), response.response().orElse(null), id); targets.forEach((interruptId, node) -> node.recordAnswer(interruptId, answer)); } } @@ -492,11 +602,26 @@ private static void recordOutcome( } /** - * Records the long-running call interrupts that {@code event} raises for {@code node}, and - * returns whether {@code event} raised any. + * Records the interrupts {@code event} raises for {@code node} (long-running calls and requests + * for input, together with their response schemas) and returns whether it raised any. */ private boolean recordRaisedInterrupts(Event event, RecoveredNode node) { - Set raised = event.longRunningToolIds().orElse(ImmutableSet.of()); + Set raised = + new LinkedHashSet<>(event.longRunningToolIds().orElse(ImmutableSet.of())); + for (FunctionCall call : event.functionCalls()) { + if (call.id().isEmpty() + || !call.name().orElse("").equals(RequestInput.FUNCTION_CALL_NAME)) { + continue; + } + String id = call.id().get(); + raised.add(id); + Schema schema = + readSchema( + call.args().map(args -> args.get(RequestInput.RESPONSE_SCHEMA_KEY)).orElse(null)); + if (schema != null) { + schemas.put(id, schema); + } + } for (String id : raised) { // An interrupt raised again waits for a new answer, as in Python. node.resolvedResponses.remove(id); diff --git a/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java index b64115532..690b681d7 100644 --- a/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java +++ b/core/src/test/java/com/google/adk/workflow/ResumeScanTest.java @@ -41,6 +41,7 @@ import com.google.genai.types.Schema; import com.google.genai.types.Type; import io.reactivex.rxjava3.core.Flowable; +import java.util.Collections; import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; @@ -57,6 +58,14 @@ public final class ResumeScanTest { .properties(ImmutableMap.of("count", Schema.builder().type(Type.Known.INTEGER).build())) .build(); + private static final Schema INTEGER_ARRAY_SCHEMA = + Schema.builder() + .type(Type.Known.ARRAY) + .items(Schema.builder().type(Type.Known.INTEGER).build()) + .build(); + + private static final Schema NULL_SCHEMA = Schema.builder().type(Type.Known.NULL).build(); + @Test public void scan_interruptRaisedBelowAChild_isAnsweredOnThatChild() { ImmutableList events = @@ -70,6 +79,20 @@ public void scan_interruptRaisedBelowAChild_isAnsweredOnThatChild() { .containsExactly("ask_1", ImmutableMap.of("v", 1)); } + @Test + public void scan_requestForInputWithoutLongRunningToolIds_isAnInterrupt() { + // Sessions recorded before requests carried long-running tool IDs hold only the call. + ImmutableList events = + ImmutableList.of( + nodeEvent("wf@1/ask@1") + .content(callContent(RequestInput.FUNCTION_CALL_NAME, "ask_1")) + .build()); + + RecoveredNode recovered = ResumeScan.scan(events, "wf@1", "inv").get("ask@1"); + + assertThat(recovered.unresolved()).containsExactly("ask_1"); + } + @Test public void scan_answerFromAnotherInvocation_resolvesNothing() { ImmutableList events = @@ -328,6 +351,118 @@ public void processRehydratedOutput_stringSchema_keepsTheText() { assertThat(ResumeScan.processRehydratedOutput(node, modelText("42"))).isEqualTo("42"); } + @Test + public void validateAnswer_answerThatFits_isReturned() { + Object answer = ResumeScan.validateAnswer(COUNT_SCHEMA, ImmutableMap.of("count", 1), "ask_1"); + + assertThat(answer).isEqualTo(ImmutableMap.of("count", 1)); + } + + @Test + public void validateAnswer_objectTheSchemaRejects_failsWithoutTheValue() { + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> + ResumeScan.validateAnswer( + COUNT_SCHEMA, ImmutableMap.of("result", "{\"count\": \"1\"}"), "ask_1")); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + assertThat(error).hasMessageThat().doesNotContain("\"1\""); + } + + @Test + public void validateAnswer_textTheSchemaRejects_failsWithoutTheValue() { + Schema booleanSchema = Schema.builder().type(Type.Known.BOOLEAN).build(); + + IllegalArgumentException error = + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.validateAnswer(booleanSchema, ImmutableMap.of("result", "yes"), "a")); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt a"); + assertThat(error).hasMessageThat().doesNotContain("yes"); + } + + @Test + public void validateAnswer_arrayOfMatchingItems_isReturned() { + Object answer = + ResumeScan.validateAnswer( + INTEGER_ARRAY_SCHEMA, ImmutableMap.of("result", ImmutableList.of(1, 2)), "a"); + + assertThat(answer).isEqualTo(ImmutableList.of(1, 2)); + } + + @Test + public void validateAnswer_arrayWithAnItemOfAnotherType_fails() { + assertThrows( + IllegalArgumentException.class, + () -> + ResumeScan.validateAnswer( + INTEGER_ARRAY_SCHEMA, ImmutableMap.of("result", ImmutableList.of(1, "2")), "a")); + } + + @Test + public void validateAnswer_numberTextForANumberSchema_isReadAsANumber() { + Schema numberSchema = Schema.builder().type(Type.Known.NUMBER).build(); + + Object answer = ResumeScan.validateAnswer(numberSchema, ImmutableMap.of("result", "2.5"), "a"); + + assertThat(answer).isEqualTo(2.5); + } + + @Test + public void validateAnswer_textForANumberSchema_fails() { + Schema numberSchema = Schema.builder().type(Type.Known.NUMBER).build(); + + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.validateAnswer(numberSchema, ImmutableMap.of("result", "many"), "a")); + } + + @Test + public void validateAnswer_nullForANullSchema_isNull() { + Object answer = + ResumeScan.validateAnswer(NULL_SCHEMA, Collections.singletonMap("result", null), "a"); + + assertThat(answer).isNull(); + } + + @Test + public void validateAnswer_valueForANullSchema_fails() { + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.validateAnswer(NULL_SCHEMA, ImmutableMap.of("result", 1), "a")); + } + + @Test + public void validateAnswer_nullForAnAnyOfThatAllowsNull_isNull() { + Schema ageOrNull = Schema.builder().anyOf(AGE_SCHEMA, NULL_SCHEMA).build(); + + Object answer = + ResumeScan.validateAnswer(ageOrNull, Collections.singletonMap("result", null), "a"); + + assertThat(answer).isNull(); + } + + @Test + public void validateAnswer_objectThatMatchesOneAnyOfAlternative_isReturned() { + Schema countOrAge = Schema.builder().anyOf(COUNT_SCHEMA, AGE_SCHEMA).build(); + + Object answer = ResumeScan.validateAnswer(countOrAge, ImmutableMap.of("age", 3), "a"); + + assertThat(answer).isEqualTo(ImmutableMap.of("age", 3)); + } + + @Test + public void validateAnswer_objectThatMatchesNoAnyOfAlternative_fails() { + Schema countOrAge = Schema.builder().anyOf(COUNT_SCHEMA, AGE_SCHEMA).build(); + + assertThrows( + IllegalArgumentException.class, + () -> ResumeScan.validateAnswer(countOrAge, ImmutableMap.of("height", 3), "a")); + } + /** * Returns an event in which the node at {@code path} waits on the long-running call {@code * interruptId}. diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowHitlTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowHitlTest.java new file mode 100644 index 000000000..5c7866522 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/WorkflowHitlTest.java @@ -0,0 +1,293 @@ +/* + * 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.emitter; +import static com.google.adk.testing.WorkflowTestNodes.workflow; +import static com.google.adk.workflow.WorkflowResumeTesting.AGE_SCHEMA; +import static com.google.adk.workflow.WorkflowResumeTesting.START; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.pathsOf; +import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.MoreCollectors.onlyElement; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.events.Event; +import com.google.adk.testing.TestNode; +import com.google.adk.workflow.WorkflowResumeTesting.FailedTurn; +import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; +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.FunctionCall; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Flowable; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for pausing a workflow with a {@link RequestInput} until the user answers it. */ +@RunWith(JUnit4.class) +public final class WorkflowHitlTest { + + @Test + public void requestInput_withoutAnInterruptId_getsOne() { + RequestInput request = RequestInput.builder().message("name?").build(); + + assertThat(request.interruptId()).isNotEmpty(); + } + + @Test + public void requestInput_toBuilder_keepsEveryProperty() { + RequestInput request = + RequestInput.builder() + .interruptId("ask_1") + .message("age?") + .payload(ImmutableMap.of("hint", "years")) + .responseSchema(AGE_SCHEMA) + .build(); + + assertThat(request.toBuilder().build()).isEqualTo(request); + } + + @Test + public void requestInput_emitted_pausesTheWorkflowWithTheRequest() { + Node ask = + ask( + "ask", + RequestInput.builder() + .interruptId("ask_1") + .message("Please provide user details.") + .payload(ImmutableMap.of("fields", "name, age")) + .build()); + Node reply = emitter("reply", "done"); + + ImmutableList events = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, reply))).turn(go()); + + Event requestEvent = requestEvents(events).stream().collect(onlyElement()); + FunctionCall call = requestEvent.functionCalls().get(0); + assertThat(call.name()).hasValue(RequestInput.FUNCTION_CALL_NAME); + assertThat(call.id()).hasValue("ask_1"); + assertThat(call.args().orElseThrow()) + .containsAtLeast( + RequestInput.MESSAGE_KEY, + "Please provide user details.", + RequestInput.PAYLOAD_KEY, + ImmutableMap.of("fields", "name, age")); + assertThat(requestEvent.longRunningToolIds()).hasValue(ImmutableSet.of("ask_1")); + // History rewriting drops an event without a role, which would orphan the answer. + assertThat(requestEvent.content().flatMap(Content::role)).hasValue("model"); + assertThat(pathsOf(events)).doesNotContain("wf@1/reply@1"); + } + + @Test + public void requestInput_withAPayloadThatHasNoJsonForm_failsTheAskingNode() { + Node ask = + ask("ask", RequestInput.builder().interruptId("ask_1").payload(new Object()).build()); + + FailedTurn turn = new WorkflowSession(workflow(new Edge(START, ask))).failingTurn(go()); + + assertThat(turn.error()).isInstanceOf(IllegalArgumentException.class); + assertThat(requestEvents(turn.events())).isEmpty(); + } + + @Test + public void resume_answeredRequest_completesTheNodeAndRunsItsSuccessor() { + Node ask = ask("ask", request("ask_1")); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + ImmutableMap details = ImmutableMap.of("name", "John", "age", 30); + + ImmutableList events = session.turn(requestAnswer("ask_1", details)); + + assertThat(replyInput(events)).isEqualTo(details); + } + + @Test + public void resume_objectAnswerWithAWrongType_failsNamingTheInterrupt() { + WorkflowSession session = sessionAwaitingAnswer(AGE_SCHEMA); + Content reply = requestAnswer("ask_1", ImmutableMap.of("age", "25")); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(reply)); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + } + + @Test + public void resume_objectAnswerWithoutARequiredProperty_failsNamingTheInterrupt() { + WorkflowSession session = sessionAwaitingAnswer(AGE_SCHEMA); + Content reply = requestAnswer("ask_1", ImmutableMap.of("wrong", 1)); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(reply)); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + } + + @Test + public void resume_objectAnswerWithAnUndeclaredKey_failsWithoutNamingTheKey() { + WorkflowSession session = sessionAwaitingAnswer(AGE_SCHEMA); + Content reply = requestAnswer("ask_1", ImmutableMap.of("age", 3, "s3cret", 1)); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(reply)); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + assertThat(error).hasMessageThat().doesNotContain("s3cret"); + } + + @Test + public void resume_objectAnswerWithoutAnOptionalProperty_leavesItOut() { + Schema schema = + Schema.builder() + .type(Type.Known.OBJECT) + .properties( + ImmutableMap.of( + "name", Schema.builder().type(Type.Known.STRING).build(), + "age", Schema.builder().type(Type.Known.INTEGER).build())) + .required("name") + .build(); + WorkflowSession session = sessionAwaitingAnswer(schema); + + ImmutableList events = + session.turn(requestAnswer("ask_1", ImmutableMap.of("name", "Ada"))); + + assertThat(replyInput(events)).isEqualTo(ImmutableMap.of("name", "Ada")); + } + + @Test + public void resume_jsonLikeTextAnswerToAStringSchema_staysText() { + WorkflowSession session = + sessionAwaitingAnswer(Schema.builder().type(Type.Known.STRING).build()); + + ImmutableList events = + session.turn(requestAnswer("ask_1", ImmutableMap.of("result", "42"))); + + assertThat(replyInput(events)).isEqualTo("42"); + } + + @Test + public void resume_textAnswerToAnAnyOfWithAString_staysText() { + Schema intOrString = + Schema.builder() + .anyOf( + Schema.builder().type(Type.Known.INTEGER).build(), + Schema.builder().type(Type.Known.STRING).build()) + .build(); + WorkflowSession session = sessionAwaitingAnswer(intOrString); + + ImmutableList events = + session.turn(requestAnswer("ask_1", ImmutableMap.of("result", "42"))); + + assertThat(replyInput(events)).isEqualTo("42"); + } + + @Test + public void resume_textAnswerToABooleanSchema_isReadAsJson() { + WorkflowSession session = + sessionAwaitingAnswer(Schema.builder().type(Type.Known.BOOLEAN).build()); + + ImmutableList events = + session.turn(requestAnswer("ask_1", ImmutableMap.of("result", "true"))); + + assertThat(replyInput(events)).isEqualTo(true); + } + + @Test + public void resume_jsonNullTextAnswerToANullableSchema_isANullAnswer() { + WorkflowSession session = + sessionAwaitingAnswer(Schema.builder().type(Type.Known.BOOLEAN).nullable(true).build()); + + ImmutableList events = + session.turn(requestAnswer("ask_1", ImmutableMap.of("result", "null"))); + + assertThat(replyInput(events)).isNull(); + } + + @Test + public void resume_nullAnswerToANonNullableSchema_failsNamingTheInterrupt() { + WorkflowSession session = sessionAwaitingAnswer(AGE_SCHEMA); + Content reply = requestAnswer("ask_1", Collections.singletonMap("result", null)); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(reply)); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + } + + @Test + public void resume_jsonNullTextAnswerToANonNullableSchema_failsNamingTheInterrupt() { + WorkflowSession session = + sessionAwaitingAnswer(Schema.builder().type(Type.Known.BOOLEAN).build()); + Content reply = requestAnswer("ask_1", ImmutableMap.of("result", "null")); + + IllegalArgumentException error = + assertThrows(IllegalArgumentException.class, () -> session.turn(reply)); + + assertThat(error).hasMessageThat().contains("Validation failed for interrupt ask_1"); + } + + /** A session paused on {@code ask_1}, whose answer must match {@code schema}. */ + private static WorkflowSession sessionAwaitingAnswer(Schema schema) { + Node ask = + ask( + "ask", + RequestInput.builder() + .interruptId("ask_1") + .message("?") + .responseSchema(schema) + .build()); + WorkflowSession session = + new WorkflowSession(workflow(new Edge(START, ask), new Edge(ask, captureInput("reply")))); + session.advance(go()); + return session; + } + + private static ImmutableList requestEvents(List events) { + return events.stream() + .filter( + event -> + event.functionCalls().stream() + .anyMatch( + call -> call.name().orElse("").equals(RequestInput.FUNCTION_CALL_NAME))) + .collect(toImmutableList()); + } + + private static Content requestAnswer(String interruptId, Map response) { + return answer(interruptId, response, RequestInput.FUNCTION_CALL_NAME); + } + + private static RequestInput request(String interruptId) { + return RequestInput.builder().interruptId(interruptId).message("need input").build(); + } + + private static Node ask(String name, RequestInput request) { + return new TestNode(name, (ctx, input) -> Flowable.just(request)); + } +}