From 8b560fcb363d95e7327624aecb651ff466904cf4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Wed, 28 Oct 2026 03:55:44 -0700 Subject: [PATCH] feat: run a tool as a workflow node PiperOrigin-RevId: 1008791546 --- .../com/google/adk/workflow/ToolNodes.java | 299 +++++++ .../google/adk/workflow/ToolNodesTest.java | 783 ++++++++++++++++++ 2 files changed, 1082 insertions(+) create mode 100644 core/src/main/java/com/google/adk/workflow/ToolNodes.java create mode 100644 core/src/test/java/com/google/adk/workflow/ToolNodesTest.java diff --git a/core/src/main/java/com/google/adk/workflow/ToolNodes.java b/core/src/main/java/com/google/adk/workflow/ToolNodes.java new file mode 100644 index 000000000..cbc220f58 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ToolNodes.java @@ -0,0 +1,299 @@ +/* + * 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.nullToEmpty; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.JsonNode; +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.annotations.Experimental; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.plugins.Plugin; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.jspecify.annotations.Nullable; + +/** + * Static factories for workflow nodes that run a {@link BaseTool}, such as an MCP tool or an agent + * wrapped as a tool. A node calls its tool between the plugins' tool callbacks and keeps only the + * tool's state and artifact changes; a tool declared to return nothing gives no output. A tool that + * asks for confirmation does not pause the workflow: the node outputs the tool's placeholder error. + */ +@Experimental +public final class ToolNodes { + + private static final TypeReference> MAP_TYPE = new TypeReference<>() {}; + + /** + * Returns a node named after {@code tool}, with the default {@link NodeConfig}. Otherwise it + * behaves as {@link #asNode(BaseTool, String, NodeConfig)} describes. + */ + public static Node asNode(BaseTool tool) { + return asNode(tool, tool.name()); + } + + /** + * Returns a node named {@code name} that runs {@code tool}, with the default {@link NodeConfig}. + * Otherwise it behaves as {@link #asNode(BaseTool, String, NodeConfig)} describes. + */ + public static Node asNode(BaseTool tool, String name) { + return asNode(tool, name, NodeConfig.builder().build()); + } + + /** + * Returns a node named {@code name} that runs {@code tool} with {@code config}; session state + * supplies any required parameter that the input lacks. The input may be a map of arguments; an + * object the JSON mapper writes as a JSON object, such as a record; a JSON object as text or as a + * {@link Content}; or null or blank text for none. A tool with one parameter, or one required + * parameter, takes any other input as that parameter; other tools fail. + */ + public static Node asNode(BaseTool tool, String name, NodeConfig config) { + return new ToolNode(tool, name, config); + } + + private static final class ToolNode extends BaseNode { + + private static final ObjectReader JSON_READER = + JsonBaseModel.getMapper() + .readerFor(Object.class) + .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + + private final BaseTool tool; + + ToolNode(BaseTool tool, String name, NodeConfig config) { + super( + name, + nullToEmpty(tool.description()), + /* rerunOnResume= */ false, + /* waitForOutput= */ false, + config, + /* inputSchema= */ null, + /* outputSchema= */ null, + /* stateSchema= */ null); + this.tool = tool; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Single.defer(() -> runTool(context, nodeInput)) + .map( + result -> + result.isEmpty() && declaresNoResult(tool) ? eventWithoutOutput(context) : result) + .toFlowable(); + } + + /** + * Returns an event without output for a tool declared to return nothing, whose result is still + * an empty map. The event records the run, so a resumed workflow does not run the tool again. + */ + private static Event eventWithoutOutput(Context context) { + InvocationContext invocationContext = context.invocationContext(); + return Event.builder() + .id(invocationContext.newUuid()) + .timestamp(invocationContext.now().toEpochMilli()) + .author("") + .build(); + } + + private Single> runTool(Context context, @Nullable Object nodeInput) { + Parameters parameters = Parameters.of(tool); + Map args = argumentsFrom(nodeInput, parameters); + for (String param : parameters.required()) { + if (!args.containsKey(param) && context.state().containsKey(param)) { + args.put(param, context.state().get(param)); + } + } + ToolContext toolContext = + ToolContext.builder(context.invocationContext()) + .functionCallId(context.invocationContext().newUuid()) + .build(); + return callTool(context.invocationContext().pluginManager(), args, toolContext) + // Kept on failure too, so the recorded deltas match the live session state. + .doOnEvent((unusedResult, unusedError) -> keepDeltas(context, toolContext)); + } + + /** Calls the tool between the plugins' tool callbacks, as a model's tool call does. */ + private Single> callTool( + Plugin plugins, Map args, ToolContext toolContext) { + Single> toolCall = + Single.defer(() -> tool.runAsync(args, toolContext)) + .onErrorResumeNext( + error -> + plugins + .onToolErrorCallback(tool, args, toolContext, error) + .switchIfEmpty(Single.error(error))); + return plugins + .beforeToolCallback(tool, args, toolContext) + .switchIfEmpty(toolCall) + .flatMap( + result -> + plugins + .afterToolCallback(tool, args, toolContext, result) + .defaultIfEmpty(result)); + } + + private Map argumentsFrom(@Nullable Object nodeInput, Parameters parameters) { + Object input = nodeInput instanceof Content content ? content.text() : nodeInput; + if (input instanceof String text) { + input = parseText(text.strip()); + } + Map args = new LinkedHashMap<>(); + if (input == null) { + return args; + } + if (input instanceof Map map) { + map.forEach((key, value) -> args.put(String.valueOf(key), value)); + return args; + } + Map properties = propertiesOf(input); + if (properties != null) { + return properties; + } + String param = parameters.forOtherInput(); + if (param == null) { + throw new IllegalArgumentException( + String.format( + "The input to tool node '%s' must be a dictionary of tool arguments, an object that" + + " maps to one, a JSON object as text, or null, but got %s.", + name(), + input instanceof String ? "text that is not a JSON object" : typeName(input))); + } + args.put(param, input); + return args; + } + + /** + * Returns the properties of an object that the JSON mapper writes as a JSON object, such as a + * record, as Python ADK reads a Pydantic model. Returns null for any other value. + */ + private static @Nullable Map propertiesOf(Object input) { + JsonNode tree; + try { + tree = JsonBaseModel.getMapper().valueToTree(input); + } catch (IllegalArgumentException e) { + return null; + } + return tree.isObject() ? JsonBaseModel.getMapper().convertValue(tree, MAP_TYPE) : null; + } + + private static String typeName(Object value) { + String name = value.getClass().getSimpleName(); + return name.isEmpty() ? value.getClass().getName() : name; + } + + private static boolean declaresNoResult(BaseTool tool) { + return tool.declaration() + .flatMap(FunctionDeclaration::response) + .flatMap(Schema::type) + .map(type -> type.knownEnum() == Type.Known.NULL) + .orElse(false); + } + + private static @Nullable Object parseText(String text) { + if (text.isEmpty()) { + return null; + } + try { + Object json = JSON_READER.readValue(text); + return json instanceof Map ? json : text; + } catch (JsonProcessingException e) { + return text; + } + } + + private static void keepDeltas(Context context, ToolContext toolContext) { + EventActions toolActions = toolContext.eventActions(); + context.eventActions().stateDelta().putAll(toolActions.stateDelta()); + context.eventActions().artifactDelta().putAll(toolActions.artifactDelta()); + } + } + + /** + * The declared parameters that a node fills itself. + * + * @param required the parameters read from state when the input lacks them + * @param forOtherInput the parameter that any other input fills, as in Python ADK: the only + * parameter, or else the only required one + */ + private record Parameters(ImmutableList required, @Nullable String forOtherInput) { + private static final Parameters NONE = new Parameters(ImmutableList.of(), null); + + /** Reads {@code tool}'s raw JSON schema, which MCP tools set, or else its typed schema. */ + static Parameters of(BaseTool tool) { + FunctionDeclaration declaration = tool.declaration().orElse(null); + if (declaration == null) { + return NONE; + } + Map rawSchema = + declaration + .parametersJsonSchema() + .map(schema -> JsonBaseModel.getMapper().convertValue(schema, MAP_TYPE)) + .orElse(ImmutableMap.of()); + ImmutableList all; + ImmutableList required; + if (!rawSchema.isEmpty()) { + all = + rawSchema.get("properties") instanceof Map properties + ? names(properties.keySet()) + : ImmutableList.of(); + required = + rawSchema.get("required") instanceof List list ? names(list) : ImmutableList.of(); + } else { + Schema schema = declaration.parameters().orElse(null); + if (schema == null) { + return NONE; + } + all = + schema + .properties() + .map(properties -> ImmutableList.copyOf(properties.keySet())) + .orElse(ImmutableList.of()); + required = schema.required().map(ImmutableList::copyOf).orElse(ImmutableList.of()); + } + if (all.size() == 1) { + return new Parameters(required, all.get(0)); + } + return new Parameters(required, required.size() == 1 ? required.get(0) : null); + } + + private static ImmutableList names(Iterable values) { + ImmutableList.Builder names = ImmutableList.builder(); + values.forEach(value -> names.add(String.valueOf(value))); + return names.build(); + } + } + + private ToolNodes() {} +} diff --git a/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java new file mode 100644 index 000000000..891d9c2dd --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java @@ -0,0 +1,783 @@ +/* + * 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.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.testing.WorkflowTestNodes.emitter; +import static com.google.adk.testing.WorkflowTestNodes.workflow; +import static com.google.adk.workflow.WorkflowResumeTesting.LONG_RUNNING_CALL; +import static com.google.adk.workflow.WorkflowResumeTesting.START; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.callContent; +import static com.google.adk.workflow.WorkflowResumeTesting.captureInput; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.newEvent; +import static com.google.adk.workflow.WorkflowResumeTesting.replyInput; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.plugins.PluginManager; +import com.google.adk.testing.TestLlm; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.adk.tools.AgentTool; +import com.google.adk.tools.Annotations; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.ToolContext; +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.common.collect.Iterables; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class ToolNodesTest { + + @Test + public void asNode_withOnlyATool_usesItsNameAndDescriptionAndDefaults() { + Node node = ToolNodes.asNode(FakeTool.echo("lookup")); + + assertThat(node.name()).isEqualTo("lookup"); + assertThat(node.description()).isEqualTo("A fake tool."); + assertThat(node.config()).isEqualTo(NodeConfig.builder().build()); + assertThat(node.inputSchema()).isEmpty(); + assertThat(node.outputSchema()).isEmpty(); + } + + @Test + public void asNode_withAName_overridesTheToolsName() { + Node node = ToolNodes.asNode(FakeTool.echo("lookup"), "custom"); + + assertThat(node.name()).isEqualTo("custom"); + assertThat(node.description()).isEqualTo("A fake tool."); + } + + @Test + public void asNode_withANameAndConfig_usesThem() { + NodeConfig config = NodeConfig.builder().timeout(Duration.ofSeconds(5)).build(); + + Node node = ToolNodes.asNode(FakeTool.echo("lookup"), "custom", config); + + assertThat(node.name()).isEqualTo("custom"); + assertThat(node.config()).isEqualTo(config); + } + + @Test + public void asNode_toolWithoutDescription_givesAnEmptyDescription() { + BaseTool tool = new BaseTool("lookup", /* description= */ null) {}; + + assertThat(ToolNodes.asNode(tool).description()).isEmpty(); + } + + @Test + public void runNode_mapInput_passesItAsTheToolsArguments() { + ImmutableMap args = ImmutableMap.of("param_a", 1, "param_b", "value"); + + assertThat(run(FakeTool.echo("mock_tool"), args)).isEqualTo(args); + } + + @Test + public void runNode_jsonObjectText_parsesItIntoArguments() { + Object output = run(FakeTool.echo("mock_tool"), " {\"param_a\": 1, \"param_b\": \"value\"} "); + + assertThat(output).isEqualTo(ImmutableMap.of("param_a", 1, "param_b", "value")); + } + + @Test + public void runNode_nullInput_callsTheToolWithNoArguments() { + assertThat(run(FakeTool.echo("mock_tool"), null)).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_blankText_callsTheToolWithNoArguments() { + assertThat(run(FakeTool.echo("mock_tool"), " \n\t")).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_jsonObjectContent_readsItsTextWithoutThoughts() { + Content content = + Content.builder() + .role("model") + .parts( + Part.builder().text("reasoning").thought(true).build(), Part.fromText("{\"x\": 1}")) + .build(); + + assertThat(run(FakeTool.echo("mock_tool"), content)).isEqualTo(ImmutableMap.of("x", 1)); + } + + @Test + public void runNode_textThatIsNotJson_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), "not a json")); + + assertThat(thrown) + .hasMessageThat() + .contains("The input to tool node 'mock_tool' must be a dictionary of tool arguments"); + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + assertThat(thrown).hasMessageThat().doesNotContain("not a json"); + } + + @Test + public void runNode_jsonArrayText_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), "[1, 2, 3]")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + assertThat(thrown).hasMessageThat().doesNotContain("[1, 2, 3]"); + } + + @Test + public void runNode_objectTextWithTrailingText_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> run(FakeTool.echo("mock_tool"), "{\"city\": \"Paris\"} and more")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + assertThat(thrown).hasMessageThat().doesNotContain("Paris"); + } + + @Test + public void runNode_recordInput_passesItsPropertiesAsArguments() { + FakeTool tool = + FakeTool.echo("weather", ImmutableList.of("city", "units"), ImmutableList.of("city")); + + assertThat(run(tool, new City("Paris", "metric"))) + .isEqualTo(ImmutableMap.of("city", "Paris", "units", "metric")); + } + + @Test + public void runNode_objectWithoutProperties_failsNamingItsClass() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), new Object() {})); + + assertThat(thrown).hasMessageThat().contains("but got com.google.adk.workflow.ToolNodesTest$"); + } + + @Test + public void runNode_nonMapValue_failsTheNode() { + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), 42)); + + assertThat(thrown).hasMessageThat().contains("but got Integer"); + } + + @Test + public void runNode_toolWithOneParameter_takesStrippedTextAsIt() { + Object output = run(FakeTool.echo("reader", "path"), " /tmp/a.txt "); + + assertThat(output).isEqualTo(ImmutableMap.of("path", "/tmp/a.txt")); + } + + @Test + public void runNode_toolWithOneParameter_takesANonMapValueAsIt() { + Object output = run(FakeTool.echo("counter", "count"), 42); + + assertThat(output).isEqualTo(ImmutableMap.of("count", 42)); + } + + @Test + public void runNode_toolWithOneParameter_takesJsonNullTextAsText() { + Object output = run(FakeTool.echo("writer", "text"), "null"); + + assertThat(output).isEqualTo(ImmutableMap.of("text", "null")); + } + + @Test + public void runNode_toolWithOneRequiredParameter_takesANonMapValueAsIt() { + FakeTool tool = + FakeTool.echo("weather", ImmutableList.of("city", "units"), ImmutableList.of("city")); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_toolWithTwoParameters_rejectsANonMapValue() { + FakeTool tool = FakeTool.echo("lookup", ImmutableList.of("a", "b"), ImmutableList.of()); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> run(tool, "scalar")); + + assertThat(thrown).hasMessageThat().contains("must be a dictionary"); + } + + @Test + public void runNode_toolWithRawJsonSchema_readsItsParameters() { + // MCP tools declare their parameters as a raw JSON schema. + FakeTool tool = + FakeTool.withRawSchema( + "weather", + ImmutableMap.of( + "type", + "object", + "properties", + ImmutableMap.of( + "city", ImmutableMap.of("type", "string"), + "units", ImmutableMap.of("type", "string")), + "required", + ImmutableList.of("city"))); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_toolWithRawAndTypedSchemas_readsTheRawOne() { + Schema typed = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "a", Schema.builder().type("STRING").build(), + "b", Schema.builder().type("STRING").build())) + .build(); + FunctionDeclaration declaration = + FunctionDeclaration.builder() + .name("weather") + .parameters(typed) + .parametersJsonSchema( + ImmutableMap.of( + "type", + "object", + "properties", + ImmutableMap.of("city", ImmutableMap.of("type", "string")))) + .build(); + FakeTool tool = new FakeTool("weather", declaration, (args, toolContext) -> args); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_requiredParametersMissing_readsOnlyThoseFromState() { + FakeTool tool = + FakeTool.echo("weather", ImmutableList.of("city", "units"), ImmutableList.of("city")); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + context.invocationContext().session().state().put("units", "fahrenheit"); + + assertThat(run(tool, context, null)).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_requiredParameterInNeitherInputNorState_isLeftOut() { + FakeTool tool = FakeTool.echo("weather", "city"); + + assertThat(run(tool, ImmutableMap.of())).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_argumentInInputAndState_usesTheInput() { + FakeTool tool = FakeTool.echo("weather", "city"); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + + assertThat(run(tool, context, ImmutableMap.of("city", "Rome"))) + .isEqualTo(ImmutableMap.of("city", "Rome")); + } + + @Test + public void runNode_functionTool_readsAMissingRequiredArgumentFromState() { + FunctionTool tool = FunctionTool.create(Functions.class, "forecast"); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + + assertThat(run(tool, context, null)).isEqualTo(ImmutableMap.of("forecast", "Sunny in Paris")); + } + + @Test + public void runNode_functionToolRequiringConfirmation_outputsItsPlaceholderWithoutRunning() { + FunctionTool tool = + FunctionTool.create(Functions.class, "forecast", /* requireConfirmation= */ true); + Context context = context(); + + Object output = run(tool, context, "Rome"); + + assertThat(output) + .isEqualTo( + ImmutableMap.of( + "error", "This tool call requires confirmation, please approve or reject.")); + assertThat(context.eventActions().requestedToolConfirmations()).isEmpty(); + } + + @Test + public void nodeRunner_toolWritesState_putsTheDeltaOnTheOutputEvent() { + FakeTool tool = + FakeTool.running( + "writer", + (args, toolContext) -> { + toolContext.state().put("color", "blue"); + return ImmutableMap.of("status", "ok"); + }); + + Event event = Iterables.getOnlyElement(runAsRoot(ToolNodes.asNode(tool))); + + assertThat(event.output()).hasValue(ImmutableMap.of("status", "ok")); + assertThat(event.actions().stateDelta()).containsExactly("color", "blue"); + } + + @Test + public void runNode_toolSavesArtifact_keepsTheArtifactDeltaOnTheNodeContext() { + FakeTool tool = + FakeTool.running( + "saver", + (args, toolContext) -> { + toolContext.saveArtifact("report.txt", Part.fromText("v1")).blockingAwait(); + return ImmutableMap.of("status", "saved"); + }); + Context context = context(); + + Object output = run(tool, context, null); + + assertThat(output).isEqualTo(ImmutableMap.of("status", "saved")); + assertThat(context.eventActions().artifactDelta()).containsExactly("report.txt", 0); + } + + @Test + public void nodeRunner_toolWritesStateThenFails_putsTheDeltaOnTheErrorEvent() { + FakeTool tool = + FakeTool.running( + "writer", + (args, toolContext) -> { + toolContext.state().put("color", "blue"); + throw new IllegalStateException("transient failure"); + }); + + Event event = Iterables.getOnlyElement(runAsRoot(ToolNodes.asNode(tool))); + + assertThat(event.errorMessage()).hasValue("transient failure"); + assertThat(event.actions().stateDelta()).containsExactly("color", "blue"); + } + + @Test + public void runNode_functionToolReturningNothing_sendsOnlyAnEventWithoutOutput() { + FunctionTool tool = FunctionTool.create(Functions.class, "record"); + + Object emitted = run(tool, "a note"); + + assertThat(emitted).isInstanceOf(Event.class); + assertThat(((Event) emitted).output()).isEmpty(); + } + + @Test + public void runNode_functionToolReturningNothingThatNeedsConfirmation_outputsItsPlaceholder() { + FunctionTool tool = + FunctionTool.create(Functions.class, "record", /* requireConfirmation= */ true); + + assertThat(run(tool, "a note")) + .isEqualTo( + ImmutableMap.of( + "error", "This tool call requires confirmation, please approve or reject.")); + } + + @Test + public void runNode_toolDeclaringAnObjectResultReturnsAnEmptyMap_outputsIt() { + FakeTool tool = + new FakeTool( + "lookup", + FunctionDeclaration.builder() + .name("lookup") + .response(Schema.builder().type("OBJECT").build()) + .build(), + (args, toolContext) -> ImmutableMap.of()); + + assertThat(run(tool, null)).isEqualTo(ImmutableMap.of()); + } + + @Test + public void resume_afterALaterInterrupt_doesNotRunAToolDeclaredToReturnNothingAgain() { + AtomicInteger calls = new AtomicInteger(); + FakeTool tool = + new FakeTool( + "notify", + FunctionDeclaration.builder() + .name("notify") + .response(Schema.builder().type("NULL").build()) + .build(), + (args, toolContext) -> { + calls.incrementAndGet(); + return ImmutableMap.of(); + }); + Node arguments = emitter("arguments", ImmutableMap.of()); + Node toolNode = ToolNodes.asNode(tool); + Node asker = new TestNode("asker", (context, input) -> Flowable.just(interrupt("ask-1"))); + WorkflowSession session = + new WorkflowSession( + workflow( + new Edge(START, arguments), + new Edge(arguments, toolNode), + new Edge(toolNode, asker), + new Edge(asker, captureInput("reply")))); + session.advance(go()); + + ImmutableList events = session.turn(answer("ask-1", ImmutableMap.of("result", "ok"))); + + assertThat(calls.get()).isEqualTo(1); + assertThat(replyInput(events)).isEqualTo("ok"); + } + + @Test + public void runNode_toolRequestsModelOnlyActions_dropsThem() { + ImmutableMap placeholder = ImmutableMap.of("error", "Needs confirmation."); + FakeTool tool = + FakeTool.running( + "handoff", + (args, toolContext) -> { + toolContext.requestConfirmation("Approve?"); + toolContext.eventActions().setTransferToAgent("elsewhere"); + toolContext.eventActions().setSkipSummarization(true); + return placeholder; + }); + Context context = context(); + + Object output = run(tool, context, null); + + assertThat(output).isEqualTo(placeholder); + EventActions actions = context.eventActions(); + assertThat(actions.requestedToolConfirmations()).isEmpty(); + assertThat(actions.transferToAgent()).isEmpty(); + assertThat(actions.skipSummarization()).isEmpty(); + } + + @Test + public void runNode_twoRuns_giveTheToolAFreshFunctionCallIdEachTime() { + List seenIds = new ArrayList<>(); + FakeTool tool = + FakeTool.running( + "record", + (args, toolContext) -> { + seenIds.add(toolContext.functionCallId().orElse("")); + return ImmutableMap.of(); + }); + Node node = ToolNodes.asNode(tool); + Context context = context(); + + Object unusedFirst = node.runNode(context, null).blockingSingle(); + Object unusedSecond = node.runNode(context, null).blockingSingle(); + + assertThat(seenIds).hasSize(2); + assertThat(seenIds.get(0)).isNotEmpty(); + assertThat(seenIds.get(1)).isNotEqualTo(seenIds.get(0)); + } + + @Test + public void runNode_agentTool_runsTheAgentOnTheRequestAndOutputsItsReply() { + TestLlm llm = createTestLlm(createTextLlmResponse("Paris is sunny.")); + AgentTool tool = AgentTool.create(createTestAgentBuilder(llm).name("forecaster").build()); + + Object output = run(tool, "Weather in Paris?"); + + assertThat(output).isEqualTo(ImmutableMap.of("result", "Paris is sunny.")); + assertThat(Iterables.getLast(llm.getLastRequest().contents()).text()) + .isEqualTo("Weather in Paris?"); + } + + @Test + public void runNode_runsThePluginsToolCallbacksAroundTheCall() { + RecordingPlugin plugin = RecordingPlugin.silent(); + + Object output = run(FakeTool.echo("lookup"), context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("city", "Paris")); + assertThat(plugin.calls) + .containsExactly("before lookup {city=Paris}", "after lookup {city=Paris}") + .inOrder(); + } + + @Test + public void runNode_pluginBeforeAnswer_replacesTheToolCall() { + List> toolCalls = new ArrayList<>(); + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + toolCalls.add(args); + return args; + }); + RecordingPlugin plugin = RecordingPlugin.answeringBefore(ImmutableMap.of("cached", true)); + + Object output = run(tool, context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("cached", true)); + assertThat(toolCalls).isEmpty(); + assertThat(plugin.calls).contains("after lookup {cached=true}"); + } + + @Test + public void runNode_pluginAfterAnswer_replacesTheOutput() { + RecordingPlugin plugin = RecordingPlugin.answeringAfter(ImmutableMap.of("redacted", true)); + + Object output = run(FakeTool.echo("lookup"), context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("redacted", true)); + } + + @Test + public void runNode_pluginErrorAnswer_becomesTheOutput() { + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + throw new IllegalStateException("no data"); + }); + RecordingPlugin plugin = RecordingPlugin.answeringError(ImmutableMap.of("error", "handled")); + + Object output = run(tool, context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("error", "handled")); + assertThat(plugin.calls) + .containsExactly( + "before lookup {city=Paris}", "error lookup no data", "after lookup {error=handled}") + .inOrder(); + } + + @Test + public void runNode_toolFailsAndNoPluginAnswers_failsTheNode() { + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + throw new IllegalStateException("no data"); + }); + RecordingPlugin plugin = RecordingPlugin.silent(); + + IllegalStateException thrown = + assertThrows( + IllegalStateException.class, + () -> run(tool, context(plugin), ImmutableMap.of("city", "Paris"))); + + assertThat(thrown).hasMessageThat().isEqualTo("no data"); + assertThat(plugin.calls).contains("error lookup no data"); + } + + private static Object run(BaseTool tool, @Nullable Object nodeInput) { + return run(tool, context(), nodeInput); + } + + private static Object run(BaseTool tool, Context context, @Nullable Object nodeInput) { + return ToolNodes.asNode(tool).runNode(context, nodeInput).blockingSingle(); + } + + private static Context context() { + return context(new PluginManager()); + } + + private static Context context(BasePlugin plugin) { + return context(new PluginManager(ImmutableList.of(plugin))); + } + + private static Context context(PluginManager pluginManager) { + InvocationContext invocationContext = + TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .pluginManager(pluginManager) + .build(); + return new CallbackContext(invocationContext, /* eventActions= */ null); + } + + /** Runs {@code node} as the root of a workflow run and returns the events it sends. */ + private static ImmutableList runAsRoot(Node node) { + List events = new ArrayList<>(); + Context unused = + NodeRunner.builder() + .node(node) + .invocationContext(TestUtils.createInvocationContext(TestUtils.createRootAgent())) + .eventSink((event, nodeFailure) -> Completable.fromAction(() -> events.add(event))) + .build() + .run(/* nodeInput= */ null) + .blockingGet(); + return ImmutableList.copyOf(events); + } + + private static Event interrupt(String id) { + return newEvent() + .content(callContent(LONG_RUNNING_CALL, id)) + .longRunningToolIds(ImmutableSet.of(id)) + .build(); + } + + private record City(String city, String units) {} + + /** Functions for the {@link FunctionTool} tests. */ + public static final class Functions { + public static ImmutableMap forecast( + @Annotations.Schema(name = "city") String city) { + return ImmutableMap.of("forecast", "Sunny in " + city); + } + + public static void record(@Annotations.Schema(name = "note") String note) {} + + private Functions() {} + } + + /** A tool with an optional declaration that runs {@code body}. */ + private static final class FakeTool extends BaseTool { + private final @Nullable FunctionDeclaration declaration; + private final BiFunction, ToolContext, Map> body; + + private FakeTool( + String name, + @Nullable FunctionDeclaration declaration, + BiFunction, ToolContext, Map> body) { + super(name, "A fake tool."); + this.declaration = declaration; + this.body = body; + } + + /** Returns a tool without a declaration that outputs its arguments. */ + static FakeTool echo(String name) { + return running(name, (args, toolContext) -> args); + } + + /** Returns a tool with one required string parameter that outputs its arguments. */ + static FakeTool echo(String name, String param) { + return echo(name, ImmutableList.of(param), ImmutableList.of(param)); + } + + /** Returns a tool that declares string {@code params} and outputs its arguments. */ + static FakeTool echo( + String name, ImmutableList params, ImmutableList required) { + Schema parameters = + Schema.builder() + .type("OBJECT") + .properties( + params.stream() + .collect( + toImmutableMap( + param -> param, param -> Schema.builder().type("STRING").build()))) + .required(required) + .build(); + return new FakeTool( + name, + FunctionDeclaration.builder().name(name).parameters(parameters).build(), + (args, toolContext) -> args); + } + + /** + * Returns a tool that declares {@code schema} as a raw JSON schema and outputs its arguments. + */ + static FakeTool withRawSchema(String name, Map schema) { + return new FakeTool( + name, + FunctionDeclaration.builder().name(name).parametersJsonSchema(schema).build(), + (args, toolContext) -> args); + } + + /** Returns a tool without a declaration that runs {@code body}. */ + static FakeTool running( + String name, BiFunction, ToolContext, Map> body) { + return new FakeTool(name, /* declaration= */ null, body); + } + + @Override + public Optional declaration() { + return Optional.ofNullable(declaration); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + return Single.fromCallable(() -> body.apply(args, toolContext)); + } + } + + /** Records the plugins' tool callbacks, answering at most one of them with a map. */ + private static final class RecordingPlugin extends BasePlugin { + final List calls = new ArrayList<>(); + private final @Nullable Map beforeAnswer; + private final @Nullable Map afterAnswer; + private final @Nullable Map errorAnswer; + + private RecordingPlugin( + @Nullable Map beforeAnswer, + @Nullable Map afterAnswer, + @Nullable Map errorAnswer) { + super("recording"); + this.beforeAnswer = beforeAnswer; + this.afterAnswer = afterAnswer; + this.errorAnswer = errorAnswer; + } + + static RecordingPlugin silent() { + return new RecordingPlugin(null, null, null); + } + + static RecordingPlugin answeringBefore(Map answer) { + return new RecordingPlugin(answer, null, null); + } + + static RecordingPlugin answeringAfter(Map answer) { + return new RecordingPlugin(null, answer, null); + } + + static RecordingPlugin answeringError(Map answer) { + return new RecordingPlugin(null, null, answer); + } + + @Override + public Maybe> beforeToolCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext) { + calls.add("before " + tool.name() + " " + toolArgs); + return beforeAnswer == null ? Maybe.empty() : Maybe.just(beforeAnswer); + } + + @Override + public Maybe> afterToolCallback( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + calls.add("after " + tool.name() + " " + result); + return afterAnswer == null ? Maybe.empty() : Maybe.just(afterAnswer); + } + + @Override + public Maybe> onToolErrorCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { + calls.add("error " + tool.name() + " " + error.getMessage()); + return errorAnswer == null ? Maybe.empty() : Maybe.just(errorAnswer); + } + } +}