From 28f2606ce05d20444bb0c5e8db37d98ec1727b59 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Mon, 19 Oct 2026 21:20:49 -0700 Subject: [PATCH] feat: run a single workflow node PiperOrigin-RevId: 1002809551 --- .../java/com/google/adk/agents/BaseAgent.java | 45 +- .../java/com/google/adk/agents/Context.java | 96 ++- .../adk/workflow/NodeExecutionState.java | 240 ++++++ .../com/google/adk/workflow/NodeRunner.java | 403 ++++++++++ .../com/google/adk/agents/BaseAgentTest.java | 49 +- .../com/google/adk/agents/ContextTest.java | 18 + .../StubNode.java => testing/TestNode.java} | 17 +- .../com/google/adk/workflow/EdgeTest.java | 5 +- .../google/adk/workflow/EdgesBuilderTest.java | 17 +- .../com/google/adk/workflow/GraphTest.java | 91 +-- .../adk/workflow/NodeExecutionStateTest.java | 54 ++ .../google/adk/workflow/NodeRunnerTest.java | 743 ++++++++++++++++++ .../com/google/adk/workflow/WorkflowTest.java | 17 +- 13 files changed, 1684 insertions(+), 111 deletions(-) create mode 100644 core/src/main/java/com/google/adk/workflow/NodeExecutionState.java create mode 100644 core/src/main/java/com/google/adk/workflow/NodeRunner.java rename core/src/test/java/com/google/adk/{workflow/StubNode.java => testing/TestNode.java} (59%) create mode 100644 core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java create mode 100644 core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index 81f735b25..4e17780f8 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -33,6 +33,7 @@ import com.google.adk.telemetry.Tracing; import com.google.adk.utils.AgentEnums.AgentOrigin; import com.google.adk.workflow.BaseNode; +import com.google.adk.workflow.NodeExecutionState; import com.google.common.collect.ImmutableList; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.errorprone.annotations.DoNotCall; @@ -269,16 +270,20 @@ public ImmutableList canonicalAfterAgentCallbacks( * instance. * * @param parentContext Parent context to copy. - * @return new context with updated branch name. + * @param keepBranch Whether to keep the parent's branch instead of appending this agent's name. + * @return the new context for this agent. */ - private InvocationContext createInvocationContext(InvocationContext parentContext) { + private InvocationContext createInvocationContext( + InvocationContext parentContext, boolean keepBranch) { InvocationContext.Builder builder = parentContext.toBuilder(); builder.agent(this); - // Check for branch to be truthy (not None, not empty string), - parentContext - .branch() - .filter(s -> !s.isEmpty()) - .ifPresent(branch -> builder.branch(branch + "." + name())); + if (!keepBranch) { + // Check for branch to be truthy (not None, not empty string), + parentContext + .branch() + .filter(s -> !s.isEmpty()) + .ifPresent(branch -> builder.branch(branch + "." + name())); + } return builder.build(); } @@ -289,24 +294,26 @@ private InvocationContext createInvocationContext(InvocationContext parentContex * @return stream of agent-generated events. */ public Flowable runAsync(InvocationContext parentContext) { - return runWithCallbacks(parentContext, this::runAsyncImpl); + return runWithCallbacks(parentContext, /* keepBranch= */ false, this::runAsyncImpl); } /** * Runs the agent with the given implementation. * * @param parentContext Parent context to inherit. + * @param keepBranch Whether to keep the parent's branch instead of appending this agent's name. * @param runImplementation The agent-specific logic to run. * @return stream of agent-generated events. */ private Flowable runWithCallbacks( InvocationContext parentContext, + boolean keepBranch, Function> runImplementation) { return Flowable.using( () -> { Context otelContext = Context.current(); return Instrumentation.recordAgentInvocation( - createInvocationContext(parentContext), this, otelContext); + createInvocationContext(parentContext, keepBranch), this, otelContext); }, agentInvocation -> { InvocationContext invocationContext = agentInvocation.getCtx(); @@ -445,7 +452,7 @@ private Maybe callCallback( * @return stream of agent-generated events. */ public Flowable runLive(InvocationContext parentContext) { - return runWithCallbacks(parentContext, this::runLiveImpl); + return runWithCallbacks(parentContext, /* keepBranch= */ false, this::runLiveImpl); } /** @@ -488,14 +495,26 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) { } /** - * Runs this agent as a workflow node through its {@link #runAsync} lifecycle, including its - * callbacks. The node input is not used. + * Runs this agent as a workflow node, on the node's branch so the nodes after it see its events. + * It runs the callbacks and {@link #runAsyncImpl} that {@link #runAsync} runs, but does not call + * an override of {@link #runAsync}. The node input is not used, and each event's author is copied + * into the node state so that the agent's events keep their own authors instead of an inherited + * author or the node's name. + * + * @throws IllegalStateException if {@code context} is not the context of a workflow node. */ @Experimental @Override public Flowable runNode( com.google.adk.agents.Context context, @Nullable Object nodeInput) { - return runAsync(context.invocationContext()); + NodeExecutionState nodeState = context.requireNodeState(); + return runWithCallbacks(context.invocationContext(), /* keepBranch= */ true, this::runAsyncImpl) + .doOnNext( + event -> { + if (!isNullOrEmpty(event.author())) { + nodeState.eventAuthor(event.author()); + } + }); } /** 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 0432c97c9..ef103f700 100644 --- a/core/src/main/java/com/google/adk/agents/Context.java +++ b/core/src/main/java/com/google/adk/agents/Context.java @@ -16,12 +16,16 @@ package com.google.adk.agents; +import com.google.adk.annotations.Experimental; import com.google.adk.artifacts.ListArtifactsResponse; import com.google.adk.events.EventActions; import com.google.adk.events.ToolConfirmation; import com.google.adk.memory.SearchMemoryResponse; import com.google.adk.sessions.State; +import com.google.adk.workflow.NodeExecutionState; +import com.google.adk.workflow.Route; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Maybe; @@ -31,12 +35,13 @@ import org.jspecify.annotations.Nullable; /** - * The context passed to callbacks and tools during an invocation. + * The context passed to callbacks, tools, and workflow nodes during an invocation. * *

Agent and model callbacks receive a {@link CallbackContext}, and tools and tool callbacks a * {@link com.google.adk.tools.ToolContext}; both extend this class, so a helper that serves both, * or a function tool's {@code toolContext} parameter, can take a {@code Context}. {@link - * #requestConfirmation()} needs a function call ID, which only a tool call has. + * #requestConfirmation()} requires a function call ID, which only a tool call has, while {@link + * #requireNodeState()} and the other workflow node methods require the context of a workflow node. */ public class Context extends ReadonlyContext { @@ -45,6 +50,7 @@ public class Context extends ReadonlyContext { private final State state; private final @Nullable String eventId; + private final @Nullable NodeExecutionState nodeState; private Optional functionCallId = Optional.empty(); private Optional toolConfirmation = Optional.empty(); @@ -59,10 +65,32 @@ public class Context extends ReadonlyContext { InvocationContext invocationContext, @Nullable EventActions eventActions, @Nullable String eventId) { + this(invocationContext, eventActions, eventId, /* nodeState= */ null); + } + + /** + * Creates the context of the workflow node run whose state is {@code nodeState}. Only the + * workflow runtime creates a node state. + */ + @Experimental + public Context(InvocationContext invocationContext, NodeExecutionState nodeState) { + this( + invocationContext, + /* eventActions= */ null, + /* eventId= */ null, + Preconditions.checkNotNull(nodeState)); + } + + private Context( + InvocationContext invocationContext, + @Nullable EventActions eventActions, + @Nullable String eventId, + @Nullable NodeExecutionState nodeState) { super(invocationContext); this.eventActions = eventActions != null ? eventActions : EventActions.builder().build(); this.state = new State(invocationContext.session().state(), this.eventActions.stateDelta()); this.eventId = eventId; + this.nodeState = nodeState; } /** Returns the delta-aware state of the current context. */ @@ -223,4 +251,68 @@ public Single searchMemory(String query) { .searchMemory( invocationContext.session().appName(), invocationContext.session().userId(), query); } + + /** + * Returns the state of the workflow node run this context belongs to. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final NodeExecutionState requireNodeState() { + Preconditions.checkState( + nodeState != null, "Only the context of a workflow node has a node state."); + return nodeState; + } + + /** + * Returns the author of the events this workflow node emits, or empty to use the node's name. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final String eventAuthor() { + return requireNodeState().eventAuthor(); + } + + /** + * Sets the author of the events this workflow node emits. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final void eventAuthor(String eventAuthor) { + requireNodeState().eventAuthor(eventAuthor); + } + + /** + * Returns this workflow node's output, or empty if it has produced none. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final Optional output() { + return requireNodeState().output(); + } + + /** + * Returns the routes this workflow node selected, or empty if it made no routing decision. An + * empty list is a decision that selects no route. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final Optional> routes() { + return requireNodeState().selectedRoutes(); + } + + /** + * Selects this workflow node's routes, which decide the outgoing edges to follow, replacing any + * earlier selection. + * + * @throws IllegalStateException if this is not the context of a workflow node. + */ + @Experimental + public final void routes(List routes) { + requireNodeState().selectRoutes(routes); + } } diff --git a/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java new file mode 100644 index 000000000..1801e52f2 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeExecutionState.java @@ -0,0 +1,240 @@ +/* + * 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.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; + +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import org.jspecify.annotations.Nullable; + +/** + * The state of one workflow node run, held by the run's {@link Context}. Only the workflow runtime + * creates it and updates the bookkeeping that records which results an event has carried. + */ +@Experimental +public final class NodeExecutionState { + + private final Node node; + private final NodeRunner.EventSink eventSink; + private final @Nullable Context parent; + private final String runId; + private final String nodePath; + private String eventAuthor; + private final Set interruptIds = new LinkedHashSet<>(); + private @Nullable Object output; + private boolean outputEmitted; + private @Nullable ImmutableList selectedRoutes; + private boolean routesEmitted; + private @Nullable Exception failure; + + private NodeExecutionState(Builder builder) { + checkState(builder.node != null, "A node state needs a node."); + checkState(builder.eventSink != null, "A node state needs an event sink."); + checkState(builder.runId != null, "A node state needs a run ID."); + this.node = builder.node; + this.eventSink = builder.eventSink; + this.parent = builder.parent; + this.runId = builder.runId; + NodeExecutionState parentState = parent == null ? null : parent.requireNodeState(); + this.nodePath = + NodeRunner.appendSegment( + parentState == null ? null : parentState.nodePath, + node.name(), + runId, + /* separator= */ '/'); + this.eventAuthor = parentState == null ? "" : parentState.eventAuthor; + } + + /** Returns a builder for the state of one node run. */ + static Builder builder() { + return new Builder(); + } + + /** Returns the node this run belongs to. */ + public Node node() { + return node; + } + + /** + * Returns the path of this run: {@code name@runId} segments joined by {@code /}, starting at the + * outermost node. + */ + public String nodePath() { + return nodePath; + } + + /** Returns the ID of this run among the node's runs, counting from {@code "1"}. */ + public String runId() { + return runId; + } + + /** Returns the context of the node that ran this node, or empty if this is a root node. */ + public Optional parent() { + return Optional.ofNullable(parent); + } + + /** + * Returns the author of the events this run emits, or empty to use the node's name. A run under a + * parent starts with the parent's author. + */ + public String eventAuthor() { + return eventAuthor; + } + + /** Sets the author of the events this run emits. */ + public void eventAuthor(String eventAuthor) { + this.eventAuthor = checkNotNull(eventAuthor); + } + + /** Returns the node's output, or empty if it has produced none. */ + public Optional output() { + return Optional.ofNullable(output); + } + + /** + * Records the node's output, which an event carries when the node finishes unless an earlier + * event already carried it. A node produces at most one output. + * + * @throws IllegalStateException if the node already has an output. + */ + public void produceOutput(Object output) { + checkNotNull(output); + checkState( + this.output == null, + "Node '%s' produced a second output; a node produces at most one output.", + node.name()); + this.output = output; + } + + /** + * Returns the routes the node selected, which decide the outgoing edges to follow, or empty if it + * made no routing decision. An empty list is a decision that selects no route. + */ + public Optional> selectedRoutes() { + return Optional.ofNullable(selectedRoutes); + } + + /** + * Selects the node's routes, replacing any earlier selection. An event carries them when the node + * finishes unless an earlier event already carried them. + */ + public void selectRoutes(List routes) { + this.selectedRoutes = ImmutableList.copyOf(routes); + this.routesEmitted = false; + } + + /** Returns the IDs of the long-running tool calls the node is waiting on. */ + public ImmutableSet interruptIds() { + return ImmutableSet.copyOf(interruptIds); + } + + NodeRunner.EventSink eventSink() { + return eventSink; + } + + /** Returns whether an event has carried the node's output. */ + boolean outputEmitted() { + return outputEmitted; + } + + void markOutputEmitted() { + outputEmitted = true; + } + + /** Returns whether an event has carried the node's current route selection. */ + boolean routesEmitted() { + return routesEmitted; + } + + void markRoutesEmitted() { + routesEmitted = true; + } + + void addInterruptIds(Collection ids) { + interruptIds.addAll(ids); + } + + /** Returns the exception the node failed with, if it failed. */ + Optional failure() { + return Optional.ofNullable(failure); + } + + void failure(Exception failure) { + this.failure = checkNotNull(failure); + } + + /** Builder for {@link NodeExecutionState}. */ + @Experimental + static final class Builder { + private @Nullable Node node; + private NodeRunner.@Nullable EventSink eventSink; + private @Nullable Context parent; + private @Nullable String runId; + + private Builder() {} + + /** Sets the node this run belongs to. */ + @CanIgnoreReturnValue + Builder node(Node node) { + this.node = checkNotNull(node); + return this; + } + + /** Sets the sink that receives the events of the run. */ + @CanIgnoreReturnValue + Builder eventSink(NodeRunner.EventSink eventSink) { + this.eventSink = checkNotNull(eventSink); + return this; + } + + /** + * Sets the context of the node that ran this node, or null for a root node. A run under a + * parent extends the parent's node path and starts with the parent's event author. + */ + @CanIgnoreReturnValue + Builder parent(@Nullable Context parent) { + this.parent = parent; + return this; + } + + /** Sets the ID of this run among the node's runs. */ + @CanIgnoreReturnValue + Builder runId(String runId) { + this.runId = checkNotNull(runId); + return this; + } + + /** + * Builds the state. + * + * @throws IllegalStateException if the node, the event sink, or the run ID is not set. + */ + NodeExecutionState build() { + return new NodeExecutionState(this); + } + } +} diff --git a/core/src/main/java/com/google/adk/workflow/NodeRunner.java b/core/src/main/java/com/google/adk/workflow/NodeRunner.java new file mode 100644 index 000000000..bd2284a56 --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/NodeRunner.java @@ -0,0 +1,403 @@ +/* + * 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.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.base.Strings.nullToEmpty; + +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.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.FinishReason; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import java.util.Map; +import org.jspecify.annotations.Nullable; + +/** + * The runner for one workflow node. It stamps every emitted event with its author, node path, and + * branch before forwarding it, and records the node's output, routes, interrupts, and failure on + * the node state of the returned context. + */ +@Experimental +final class NodeRunner { + + /** Sink that receives the events of a workflow run. */ + @Experimental + interface EventSink { + + /** + * Sends {@code event} and returns a {@link Completable} that completes once the event is + * processed. The Completable must not fail; a consumer that cannot process an event stops the + * run by canceling it. + */ + Completable send(Event event); + } + + private final Node node; + private final InvocationContext invocationContext; + private final EventSink eventSink; + private final @Nullable Context parent; + private final String runId; + private final boolean useSubBranch; + private final @Nullable String overrideBranch; + + private NodeRunner(Builder builder) { + checkState(builder.node != null, "A node runner needs a node."); + if (builder.parent != null) { + checkState( + builder.invocationContext == null && builder.eventSink == null, + "A node runner with a parent takes its invocation context and event sink from it."); + this.invocationContext = builder.parent.invocationContext(); + this.eventSink = builder.parent.requireNodeState().eventSink(); + } else { + checkState( + builder.invocationContext != null && builder.eventSink != null, + "A node runner without a parent needs an invocation context and an event sink."); + this.invocationContext = builder.invocationContext; + this.eventSink = builder.eventSink; + } + this.node = builder.node; + this.parent = builder.parent; + this.runId = builder.runId; + this.useSubBranch = builder.useSubBranch; + this.overrideBranch = builder.overrideBranch; + } + + /** Returns a builder for a runner of one node. */ + static Builder builder() { + return new Builder(); + } + + /** + * Runs the node on {@code nodeInput} and returns its context when the node finishes. An exception + * from the node is recorded on the context's node state rather than signaled, while an {@code + * Error} propagates. + */ + Single run(@Nullable Object nodeInput) { + return Single.defer( + () -> { + Context context = + new Context( + childInvocationContext(), + NodeExecutionState.builder() + .node(node) + .eventSink(eventSink) + .parent(parent) + .runId(runId) + .build()); + return new Activation(context).run(nodeInput).toSingleDefault(context); + }); + } + + private InvocationContext childInvocationContext() { + if (useSubBranch) { + String base = + overrideBranch != null ? overrideBranch : invocationContext.branch().orElse(null); + return invocationContext.toBuilder() + .branch(appendSegment(base, node.name(), runId, /* separator= */ '.')) + .build(); + } + if (overrideBranch != null) { + return invocationContext.toBuilder().branch(overrideBranch).build(); + } + return invocationContext; + } + + /** + * One run of the node, during which the node is asked for each event only after its previous + * event is processed. Events from nested nodes pass through without reapplying their routing or + * transfers. When the node finishes, a final event carries any output, routes, or deltas not + * carried by an earlier event. + */ + private final class Activation { + private final Context context; + private final NodeExecutionState nodeState; + + Activation(Context context) { + this.context = context; + this.nodeState = context.requireNodeState(); + } + + Completable run(@Nullable Object nodeInput) { + return asBaseNode(node) + .run(context, nodeInput) + // Avoid concatMapCompletable: its prefetch would let the node run ahead of the sink. + .flatMapCompletable(this::dispatch, /* delayErrors= */ true, /* maxConcurrency= */ 1) + .andThen(Completable.defer(this::flushPending)) + .onErrorResumeNext(this::onError); + } + + private Completable onError(Throwable error) { + // Only an Exception fails the node; an Error is fatal and propagates. + return error instanceof Exception exception ? fail(exception) : Completable.error(error); + } + + private Completable dispatch(Event event) { + track(event); + return send(event); + } + + /** Records the output, interrupts, and routing carried by {@code event} on the node state. */ + private void track(Event event) { + if (event.output().isPresent()) { + nodeState.produceOutput(event.output().get()); + } else if (isMessageAsOutput(event)) { + nodeState.produceOutput(event.content().get()); + } + event.longRunningToolIds().ifPresent(nodeState::addInterruptIds); + // An event of a nested node was already acted on at that node's level. + if (isNullOrEmpty(event.author()) || event.author().equals(node.name())) { + event + .actions() + .route() + .ifPresent( + route -> { + nodeState.selectRoutes(route); + nodeState.markRoutesEmitted(); + }); + event.actions().transferToAgent().ifPresent(context.eventActions()::setTransferToAgent); + } + } + + /** + * Stamps {@code event}, moves the pending deltas onto it unless it is partial, and sends it. + */ + private Completable send(Event event) { + boolean messageAsOutput = isMessageAsOutput(event); + // A message-as-output event's content is its output, so drop the repeated output field. + Event stamped = stamp(messageAsOutput ? event.toBuilder().output(null).build() : event); + Event outgoing = stamped.partial().orElse(false) ? stamped : withPendingDeltas(stamped); + return eventSink + .send(outgoing) + .doOnComplete( + () -> { + if (outgoing.output().isPresent() || messageAsOutput) { + nodeState.markOutputEmitted(); + } + }); + } + + /** Sends a final event carrying any output, routes, or deltas that no earlier event carried. */ + private Completable flushPending() { + boolean pendingOutput = nodeState.output().isPresent() && !nodeState.outputEmitted(); + boolean pendingRoute = nodeState.selectedRoutes().isPresent() && !nodeState.routesEmitted(); + EventActions actions = context.eventActions(); + boolean pendingDeltas = !actions.stateDelta().isEmpty() || !actions.artifactDelta().isEmpty(); + if (!pendingOutput && !pendingRoute && !pendingDeltas) { + return Completable.complete(); + } + Event event = + withPendingDeltas( + stamp( + newEvent() + .output(pendingOutput ? nodeState.output().get() : null) + .actions( + EventActions.builder() + .route(pendingRoute ? nodeState.selectedRoutes().get() : null) + .build()) + .build())); + return eventSink + .send(event) + .doOnComplete( + () -> { + if (pendingOutput) { + nodeState.markOutputEmitted(); + } + if (pendingRoute) { + nodeState.markRoutesEmitted(); + } + }); + } + + /** + * Sends an error event for {@code failure} with any pending deltas to preserve state changes + * made before the failure, then records the failure on the node state. + */ + private Completable fail(Exception failure) { + Event errorEvent = + withPendingDeltas( + stamp( + newEvent() + .errorCode(new FinishReason(errorTypeName(failure))) + .errorMessage(nullToEmpty(failure.getMessage())) + .build())); + return eventSink.send(errorEvent).doOnComplete(() -> nodeState.failure(failure)); + } + + /** + * Moves the context's pending deltas onto a copy of {@code event} so each is sent only once. + */ + private Event withPendingDeltas(Event event) { + Map stateDelta = context.eventActions().stateDelta(); + Map artifactDelta = context.eventActions().artifactDelta(); + if (stateDelta.isEmpty() && artifactDelta.isEmpty()) { + return event; + } + EventActions actions = event.actions().toBuilder().build(); + actions.stateDelta().putAll(stateDelta); + actions.artifactDelta().putAll(artifactDelta); + stateDelta.clear(); + artifactDelta.clear(); + return event.toBuilder().actions(actions).build(); + } + + /** Attributes {@code event} to its author and to this run of the node. */ + private Event stamp(Event event) { + InvocationContext eventContext = context.invocationContext(); + String author = nodeState.eventAuthor(); + NodeInfo nodeInfo = event.nodeInfo().orElseGet(() -> NodeInfo.builder().build()); + boolean carriesOutput = event.output().isPresent() || isMessageAsOutput(event); + return event.toBuilder() + .author(author.isEmpty() ? node.name() : author) + .invocationId(eventContext.invocationId()) + .branch(event.branch().orElse(eventContext.branch().orElse(null))) + .nodeInfo( + nodeInfo.toBuilder() + .path(nodeState.nodePath()) + .outputFor( + carriesOutput ? ImmutableList.of(nodeState.nodePath()) : nodeInfo.outputFor()) + .build()) + .build(); + } + + private Event.Builder newEvent() { + InvocationContext eventContext = context.invocationContext(); + return Event.builder() + .id(eventContext.newUuid()) + .timestamp(eventContext.now().toEpochMilli()); + } + } + + /** Appends a {@code name@runId} segment to {@code base}, joined by {@code separator}. */ + static String appendSegment(@Nullable String base, String name, String runId, char separator) { + String segment = name + "@" + runId; + return isNullOrEmpty(base) ? segment : base + separator + segment; + } + + /** Returns whether {@code event} is a final event whose content is the node's output. */ + private static boolean isMessageAsOutput(Event event) { + return !event.partial().orElse(false) + && event.content().isPresent() + && event.nodeInfo().map(NodeInfo::messageAsOutput).orElse(false); + } + + /** Returns the simple name of the failure's type, or its full name for an anonymous class. */ + private static String errorTypeName(Exception failure) { + String name = failure.getClass().getSimpleName(); + return name.isEmpty() ? failure.getClass().getName() : name; + } + + /** Adapts any node to a {@link BaseNode}, whose {@link BaseNode#run} turns values into events. */ + private static BaseNode asBaseNode(Node node) { + if (node instanceof BaseNode baseNode) { + return baseNode; + } + return new BaseNode(node.name(), node.description()) { + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return node.runNode(context, nodeInput); + } + }; + } + + /** + * Builder for {@link NodeRunner}. A runner at the root of a run takes an invocation context and + * an event sink, while a runner under a parent node takes both from the parent's context. + */ + @Experimental + static final class Builder { + private @Nullable Node node; + private @Nullable InvocationContext invocationContext; + private @Nullable EventSink eventSink; + private @Nullable Context parent; + private String runId = "1"; + private boolean useSubBranch; + private @Nullable String overrideBranch; + + private Builder() {} + + /** Sets the node to run. */ + @CanIgnoreReturnValue + Builder node(Node node) { + this.node = checkNotNull(node); + return this; + } + + /** Sets the invocation context of a run at the root. */ + @CanIgnoreReturnValue + Builder invocationContext(InvocationContext invocationContext) { + this.invocationContext = checkNotNull(invocationContext); + return this; + } + + /** Sets the sink that receives the events of a run at the root. */ + @CanIgnoreReturnValue + Builder eventSink(EventSink eventSink) { + this.eventSink = checkNotNull(eventSink); + return this; + } + + /** Sets the context of the node that runs this node. */ + @CanIgnoreReturnValue + Builder parent(Context parent) { + this.parent = checkNotNull(parent); + return this; + } + + /** Sets the ID of this run among the node's runs, which defaults to {@code "1"}. */ + @CanIgnoreReturnValue + Builder runId(String runId) { + this.runId = checkNotNull(runId); + return this; + } + + /** + * Sets whether the node runs on a sub-branch of its base branch, as a node fanned out from its + * parent does. Branches isolate one parallel branch's events from a sibling's model history. + */ + @CanIgnoreReturnValue + Builder useSubBranch(boolean useSubBranch) { + this.useSubBranch = useSubBranch; + return this; + } + + /** Sets the branch that replaces the inherited one, or null to keep the inherited branch. */ + @CanIgnoreReturnValue + Builder overrideBranch(@Nullable String overrideBranch) { + this.overrideBranch = overrideBranch; + return this; + } + + /** + * Builds the runner. + * + * @throws IllegalStateException if no node is set, if neither a parent nor both an invocation + * context and an event sink are set, or if a parent is set together with either of them. + */ + NodeRunner build() { + return new NodeRunner(this); + } + } +} diff --git a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java index 2792fbb76..653749023 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -26,7 +26,6 @@ import com.google.adk.testing.TestBaseAgent; import com.google.adk.testing.TestCallback; import com.google.adk.testing.TestUtils; -import com.google.adk.workflow.Node; import com.google.common.collect.ImmutableList; import com.google.genai.types.Content; import com.google.genai.types.Part; @@ -407,34 +406,6 @@ public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback() assertThat(afterCallback2.wasCalled()).isFalse(); } - @Test - public void runNode_runsAgentWithItsCallbacks() { - var runAsyncImpl = TestCallback.returningEmpty(); - Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output")); - Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output")); - var beforeCallback = TestCallback.returningEmpty(); - var afterCallback = TestCallback.returning(afterCallbackContent); - TestBaseAgent agent = - new TestBaseAgent( - TEST_AGENT_NAME, - TEST_AGENT_DESCRIPTION, - ImmutableList.of(beforeCallback.asBeforeAgentCallback()), - ImmutableList.of(afterCallback.asAfterAgentCallback()), - runAsyncImpl.asRunAsyncImplSupplier(runAsyncImplContent)); - Context context = - new CallbackContext(TestUtils.createInvocationContext(agent), /* eventActions= */ null); - Node node = agent; - - List results = - node.runNode(context, /* nodeInput= */ "unused").cast(Event.class).toList().blockingGet(); - - assertThat(results).hasSize(2); - assertThat(results.get(0).content()).hasValue(runAsyncImplContent); - assertThat(results.get(1).content()).hasValue(afterCallbackContent); - assertThat(beforeCallback.wasCalled()).isTrue(); - assertThat(runAsyncImpl.wasCalled()).isTrue(); - } - @Test public void canonicalCallbacks_returnsEmptyListWhenNull() { TestBaseAgent agent = @@ -480,6 +451,26 @@ public void runLive_invokesRunLiveImpl() { assertThat(runLiveCallback.wasCalled()).isTrue(); } + @Test + public void runNode_withoutANodeContext_throwsBeforeRunningTheAgent() { + var beforeCallback = TestCallback.returningEmpty(); + var runAsyncImpl = TestCallback.returningEmpty(); + TestBaseAgent agent = + new TestBaseAgent( + TEST_AGENT_NAME, + TEST_AGENT_DESCRIPTION, + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + /* afterAgentCallbacks= */ ImmutableList.of(), + runAsyncImpl.asRunAsyncImplSupplier(Content.fromParts(Part.fromText("main_output")))); + CallbackContext context = + new CallbackContext(TestUtils.createInvocationContext(agent), /* eventActions= */ null); + + assertThrows(IllegalStateException.class, () -> agent.runNode(context, /* nodeInput= */ null)); + + assertThat(beforeCallback.wasCalled()).isFalse(); + assertThat(runAsyncImpl.wasCalled()).isFalse(); + } + @Test public void runLive_beforeAgentCallbackReturnsContent_endsInvocationAndSkipsRunLiveImplAndAfterCallback() { diff --git a/core/src/test/java/com/google/adk/agents/ContextTest.java b/core/src/test/java/com/google/adk/agents/ContextTest.java index 6ff2b4365..c3a75ca60 100644 --- a/core/src/test/java/com/google/adk/agents/ContextTest.java +++ b/core/src/test/java/com/google/adk/agents/ContextTest.java @@ -76,6 +76,24 @@ public void requestConfirmation_onCallbackContext_throwsIllegalStateException() assertThrows(IllegalStateException.class, callbackContext::requestConfirmation); } + @Test + public void requireNodeState_onCallbackContext_throwsIllegalStateException() { + CallbackContext context = new CallbackContext(newInvocationContext(), null); + + assertThrows(IllegalStateException.class, context::requireNodeState); + } + + @Test + public void nodeMethods_onCallbackContext_throwIllegalStateException() { + CallbackContext context = new CallbackContext(newInvocationContext(), null); + + assertThrows(IllegalStateException.class, context::eventAuthor); + assertThrows(IllegalStateException.class, () -> context.eventAuthor("author")); + assertThrows(IllegalStateException.class, context::output); + assertThrows(IllegalStateException.class, context::routes); + assertThrows(IllegalStateException.class, () -> context.routes(ImmutableList.of())); + } + @Test public void searchMemory_onCallbackContext_findsRememberedEvents() { InvocationContext invocationContext = newInvocationContext(); diff --git a/core/src/test/java/com/google/adk/workflow/StubNode.java b/core/src/test/java/com/google/adk/testing/TestNode.java similarity index 59% rename from core/src/test/java/com/google/adk/workflow/StubNode.java rename to core/src/test/java/com/google/adk/testing/TestNode.java index 80a7e0d82..9bf8c9979 100644 --- a/core/src/test/java/com/google/adk/workflow/StubNode.java +++ b/core/src/test/java/com/google/adk/testing/TestNode.java @@ -14,16 +14,25 @@ * limitations under the License. */ -package com.google.adk.workflow; +package com.google.adk.testing; import com.google.adk.agents.Context; +import com.google.adk.workflow.Node; import io.reactivex.rxjava3.core.Flowable; +import java.util.function.BiFunction; import org.jspecify.annotations.Nullable; -/** A node that emits nothing, for tests that only need a named node in a graph. */ -record StubNode(String name) implements Node { +/** A workflow node whose run is given by a function, so a test can declare its node inline. */ +public record TestNode(String name, BiFunction> body) + implements Node { + + /** Returns a node named {@code name} that emits nothing. */ + public static TestNode of(String name) { + return new TestNode(name, (context, nodeInput) -> Flowable.empty()); + } + @Override public Flowable runNode(Context context, @Nullable Object nodeInput) { - return Flowable.empty(); + return body.apply(context, nodeInput); } } diff --git a/core/src/test/java/com/google/adk/workflow/EdgeTest.java b/core/src/test/java/com/google/adk/workflow/EdgeTest.java index 9fd33b088..2dd0bf85a 100644 --- a/core/src/test/java/com/google/adk/workflow/EdgeTest.java +++ b/core/src/test/java/com/google/adk/workflow/EdgeTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import org.junit.Test; import org.junit.runner.RunWith; @@ -27,8 +28,8 @@ @RunWith(JUnit4.class) public final class EdgeTest { - private final Node from = new StubNode("from"); - private final Node to = new StubNode("to"); + private final Node from = TestNode.of("from"); + private final Node to = TestNode.of("to"); @Test public void constructor_withoutRoutes_createsAnUnconditionalEdge() { diff --git a/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java index 248dcf4f7..dd1842d4c 100644 --- a/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java +++ b/core/src/test/java/com/google/adk/workflow/EdgesBuilderTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import java.util.Arrays; import java.util.concurrent.atomic.AtomicReference; @@ -30,10 +31,10 @@ public final class EdgesBuilderTest { private final EdgesBuilder builder = new EdgesBuilder(); - private final Node a = new StubNode("a"); - private final Node b = new StubNode("b"); - private final Node c = new StubNode("c"); - private final Node d = new StubNode("d"); + private final Node a = TestNode.of("a"); + private final Node b = TestNode.of("b"); + private final Node c = TestNode.of("c"); + private final Node d = TestNode.of("d"); private final JoinNode join = JoinNode.builder().name("join").build(); @Test @@ -334,10 +335,10 @@ public void anyOf_everyOverload_dropsRepeatedRoutes() { @Test public void build_documentedExample_buildsItsGraph() { - Node classify = new StubNode("classifier"); - Node category = new StubNode("category"); - Node billing = new StubNode("billing"); - Node summarize = new StubNode("summary"); + Node classify = TestNode.of("classifier"); + Node category = TestNode.of("category"); + Node billing = TestNode.of("billing"); + Node summarize = TestNode.of("summary"); builder .fromStart() diff --git a/core/src/test/java/com/google/adk/workflow/GraphTest.java b/core/src/test/java/com/google/adk/workflow/GraphTest.java index 6b7be3785..7a930c4b6 100644 --- a/core/src/test/java/com/google/adk/workflow/GraphTest.java +++ b/core/src/test/java/com/google/adk/workflow/GraphTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.adk.testing.TestNode; import com.google.common.collect.ImmutableList; import org.junit.Test; import org.junit.runner.RunWith; @@ -32,9 +33,9 @@ public final class GraphTest { @Test public void of_listsEachNodeOnceInTheOrderTheEdgesFirstMentionIt() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); - Node c = new StubNode("c"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); + Node c = TestNode.of("c"); ImmutableList edges = ImmutableList.of( new Edge(Start.INSTANCE, a), new Edge(a, c), new Edge(a, b), new Edge(b, c)); @@ -47,21 +48,21 @@ public void of_listsEachNodeOnceInTheOrderTheEdgesFirstMentionIt() { @Test public void terminalNodeNames_areTheNodesWithNoOutgoingEdge() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c")))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c")))); assertThat(graph.terminalNodeNames()).containsExactly("b", "c"); } @Test public void node_returnsTheNodeWithThatName() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, a))); assertThat(graph.node("a")).isSameInstanceAs(a); @@ -69,7 +70,7 @@ public void node_returnsTheNodeWithThatName() { @Test public void node_rejectsAnUnknownName() { - Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, new StubNode("a")))); + Graph graph = Graph.of(ImmutableList.of(new Edge(Start.INSTANCE, TestNode.of("a")))); IllegalArgumentException e = assertThrows(IllegalArgumentException.class, () -> graph.node("missing")); @@ -79,26 +80,26 @@ public void node_rejectsAnUnknownName() { @Test public void nodesTriggeredBy_followsOnlyTheEdgeWithTheEmittedRoute() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("no_node"), NO))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("no_node"), NO))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(YES))).containsExactly("yes_node"); } @Test public void nodesTriggeredBy_followsTheEdgeOfEachEmittedRoute() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("no_node"), NO))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("no_node"), NO))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO, YES))) .containsExactly("yes_node", "no_node") @@ -107,25 +108,25 @@ public void nodesTriggeredBy_followsTheEdgeOfEachEmittedRoute() { @Test public void nodesTriggeredBy_followsAnEdgeOnAnyOfItsRoutes() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("either"), ImmutableList.of(YES, NO)))); + new Edge(router, TestNode.of("either"), ImmutableList.of(YES, NO)))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))).containsExactly("either"); } @Test public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("fallback"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("fallback"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(new Route.Tag("other")))) .containsExactly("fallback"); @@ -133,13 +134,13 @@ public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteMatches() { @Test public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteIsEmitted() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("yes_node"), YES), - new Edge(router, new StubNode("fallback"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("yes_node"), YES), + new Edge(router, TestNode.of("fallback"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of())).containsExactly("fallback"); } @@ -147,14 +148,14 @@ public void nodesTriggeredBy_followsTheDefaultEdgeWhenNoRouteIsEmitted() { /** Mirrors adk-python's test_get_next_pending_nodes_with_default_route_fan_out. */ @Test public void nodesTriggeredBy_followsEveryDefaultEdgeWhenNoRouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("named"), YES), - new Edge(router, new StubNode("first"), Route.Default.INSTANCE), - new Edge(router, new StubNode("second"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("named"), YES), + new Edge(router, TestNode.of("first"), Route.Default.INSTANCE), + new Edge(router, TestNode.of("second"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))) .containsExactly("first", "second") @@ -163,42 +164,42 @@ public void nodesTriggeredBy_followsEveryDefaultEdgeWhenNoRouteMatches() { @Test public void nodesTriggeredBy_skipsEveryDefaultEdgeWhenARouteMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, router), - new Edge(router, new StubNode("named"), YES), - new Edge(router, new StubNode("first"), Route.Default.INSTANCE), - new Edge(router, new StubNode("second"), Route.Default.INSTANCE))); + new Edge(router, TestNode.of("named"), YES), + new Edge(router, TestNode.of("first"), Route.Default.INSTANCE), + new Edge(router, TestNode.of("second"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(YES))).containsExactly("named"); } @Test public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideTheDefaultEdge() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c"), YES), - new Edge(a, new StubNode("d"), Route.Default.INSTANCE))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c"), YES), + new Edge(a, TestNode.of("d"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("a", ImmutableList.of())).containsExactly("b", "d").inOrder(); } @Test public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideAMatchedRoute() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Graph graph = Graph.of( ImmutableList.of( new Edge(Start.INSTANCE, a), - new Edge(a, new StubNode("b")), - new Edge(a, new StubNode("c"), YES), - new Edge(a, new StubNode("d"), Route.Default.INSTANCE))); + new Edge(a, TestNode.of("b")), + new Edge(a, TestNode.of("c"), YES), + new Edge(a, TestNode.of("d"), Route.Default.INSTANCE))); assertThat(graph.nodesTriggeredBy("a", ImmutableList.of(YES))) .containsExactly("b", "c") @@ -207,8 +208,8 @@ public void nodesTriggeredBy_followsUnconditionalEdgesAlongsideAMatchedRoute() { @Test public void nodesTriggeredBy_followsARoutedEdgeBackToAnEarlierNode() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); Graph graph = Graph.of( ImmutableList.of(new Edge(Start.INSTANCE, a), new Edge(a, b), new Edge(b, a, YES))); @@ -218,20 +219,20 @@ public void nodesTriggeredBy_followsARoutedEdgeBackToAnEarlierNode() { @Test public void nodesTriggeredBy_returnsNothingWhenNoRoutedEdgeMatches() { - Node router = new StubNode("router"); + Node router = TestNode.of("router"); Graph graph = Graph.of( ImmutableList.of( - new Edge(Start.INSTANCE, router), new Edge(router, new StubNode("yes_node"), YES))); + new Edge(Start.INSTANCE, router), new Edge(router, TestNode.of("yes_node"), YES))); assertThat(graph.nodesTriggeredBy("router", ImmutableList.of(NO))).isEmpty(); } @Test public void predecessorsOf_returnsTheSourcesOfTheIncomingEdges() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); - Node join = new StubNode("join"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); + Node join = TestNode.of("join"); Graph graph = Graph.of( ImmutableList.of( diff --git a/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java new file mode 100644 index 000000000..9ad3e60c9 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/NodeExecutionStateTest.java @@ -0,0 +1,54 @@ +/* + * 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 org.junit.Assert.assertThrows; + +import com.google.adk.testing.TestNode; +import io.reactivex.rxjava3.core.Completable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class NodeExecutionStateTest { + + private static final NodeRunner.EventSink SINK = event -> Completable.complete(); + + @Test + public void build_withoutANode_throwsIllegalStateException() { + NodeExecutionState.Builder builder = NodeExecutionState.builder().eventSink(SINK).runId("1"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withoutAnEventSink_throwsIllegalStateException() { + NodeExecutionState.Builder builder = + NodeExecutionState.builder().node(TestNode.of("worker")).runId("1"); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_withoutARunId_throwsIllegalStateException() { + NodeExecutionState.Builder builder = + NodeExecutionState.builder().node(TestNode.of("worker")).eventSink(SINK); + + assertThrows(IllegalStateException.class, builder::build); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java new file mode 100644 index 000000000..5f04c5a07 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/NodeRunnerTest.java @@ -0,0 +1,743 @@ +/* + * 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.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +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.testing.TestBaseAgent; +import com.google.adk.testing.TestCallback; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.observers.TestObserver; +import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.util.ArrayList; +import java.util.List; +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 NodeRunnerTest { + + private static final Content MESSAGE = Content.fromParts(Part.fromText("hi")); + private static final FinishReason ILLEGAL_STATE = new FinishReason("IllegalStateException"); + + private final List sentEvents = new ArrayList<>(); + private final NodeRunner.EventSink recordingSink = + event -> Completable.fromAction(() -> sentEvents.add(event)); + private final List acknowledgements = new ArrayList<>(); + private final NodeRunner.EventSink acknowledgingSink = + event -> { + sentEvents.add(event); + CompletableSubject acknowledgement = CompletableSubject.create(); + acknowledgements.add(acknowledgement); + return acknowledgement; + }; + + @Test + public void run_stampsTheEventWithTheNodeNamePathAndInvocation() { + runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just("hi"))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.author()).isEqualTo("greet"); + assertThat(event.invocationId()).isEqualTo("invocationId"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("greet@1"); + } + + @Test + public void run_yieldedValue_becomesTheEventAndNodeStateOutput() { + Context context = runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just("hi"))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.output()).hasValue("hi"); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("greet@1")); + assertThat(context.requireNodeState().output()).hasValue("hi"); + } + + @Test + public void run_passesTheInputToTheNode() { + List inputs = new ArrayList<>(); + Node node = + new TestNode( + "echo", + (ctx, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + + rootRunner(node, recordingSink).run("question").blockingSubscribe(); + + assertThat(inputs).containsExactly("question"); + } + + @Test + public void run_withoutABranchDirective_keepsTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ false, + /* overrideBranch= */ null); + + assertThat(context.branch()).hasValue("root"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("root"); + } + + @Test + public void run_onASubBranch_derivesItFromTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ true, + /* overrideBranch= */ null); + + assertThat(context.branch()).hasValue("root.greet@1"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("root.greet@1"); + } + + @Test + public void run_withAnOverrideBranch_replacesTheParentBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ false, + /* overrideBranch= */ "merged"); + + assertThat(context.branch()).hasValue("merged"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("merged"); + } + + @Test + public void run_onASubBranchWithAnOverrideBranch_derivesItFromTheOverrideBranch() { + Context context = + runUnderParent( + new TestNode("greet", (ctx, input) -> Flowable.just("hi")), + /* useSubBranch= */ true, + /* overrideBranch= */ "merged"); + + assertThat(context.branch()).hasValue("merged.greet@1"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("merged.greet@1"); + } + + @Test + public void run_underAParent_extendsItsPathWithTheRunId() { + Context context = runWorkerUnder(parentContext()); + + assertThat(getOnlyElement(sentEvents).nodeInfo().map(NodeInfo::path)) + .hasValue("root@1/worker@2"); + assertThat(context.requireNodeState().runId()).isEqualTo("2"); + } + + @Test + public void run_underAParent_linksTheContextToTheParent() { + Context parent = parentContext(); + + Context context = runWorkerUnder(parent); + + assertThat(context.requireNodeState().parent()).hasValue(parent); + } + + @Test + public void run_underAParent_inheritsItsEventAuthor() { + Context parent = parentContext(); + parent.requireNodeState().eventAuthor("custom-agent"); + + runWorkerUnder(parent); + + assertThat(getOnlyElement(sentEvents).author()).isEqualTo("custom-agent"); + } + + @Test + public void run_messageAsOutputEvent_makesItsContentTheOutputWithoutAnotherEvent() { + Context context = + runAtRoot( + new TestNode( + "greet", (ctx, input) -> Flowable.just(messageAsOutput(/* output= */ null)))); + + Event event = getOnlyElement(sentEvents); + assertThat(context.requireNodeState().output()).hasValue(MESSAGE); + assertThat(event.content()).hasValue(MESSAGE); + assertThat(event.output()).isEmpty(); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("greet@1")); + } + + @Test + public void run_messageAsOutputEventWithAnOutputField_dropsTheField() { + runAtRoot(new TestNode("greet", (ctx, input) -> Flowable.just(messageAsOutput("hi")))); + + Event event = getOnlyElement(sentEvents); + assertThat(event.content()).hasValue(MESSAGE); + assertThat(event.output()).isEmpty(); + } + + @Test + public void run_partialMessageAsOutputEvent_isNotTheOutput() { + Context context = + runAtRoot( + new TestNode( + "greet", + (ctx, input) -> + Flowable.just( + newEvent() + .partial(true) + .content(MESSAGE) + .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) + .build()))); + + assertThat(context.requireNodeState().output()).isEmpty(); + } + + @Test + public void run_eventWithoutAnOutput_isNotMarkedAsAnyNodesOutput() { + runAtRoot( + new TestNode("greet", (ctx, input) -> Flowable.just(newEvent().content(MESSAGE).build()))); + + assertThat(getOnlyElement(sentEvents).nodeInfo().map(NodeInfo::outputFor)) + .hasValue(ImmutableList.of()); + } + + @Test + public void run_routeOnTheNodesEvent_isRecordedWithoutAnotherEvent() { + ImmutableList route = ImmutableList.of(new Route.Tag("next")); + + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent().actions(EventActions.builder().route(route).build()).build()))); + + assertThat(context.requireNodeState().selectedRoutes()).hasValue(route); + assertThat(getOnlyElement(sentEvents).actions().route()).hasValue(route); + } + + @Test + public void run_routeSelectedOnTheNodeState_isSentOnALastEvent() { + ImmutableList route = ImmutableList.of(new Route.Tag("next")); + + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> { + ctx.requireNodeState().selectRoutes(route); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().route()).hasValue(route); + assertThat(context.requireNodeState().routesEmitted()).isTrue(); + } + + @Test + public void run_nodeUsingItsContext_setsAndReadsItsNodeState() { + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> { + ctx.eventAuthor("custom-author"); + ctx.routes(ImmutableList.of(new Route.Tag("next"))); + return Flowable.just("done"); + })); + + NodeExecutionState nodeState = context.requireNodeState(); + assertThat(nodeState.eventAuthor()).isEqualTo("custom-author"); + assertThat(nodeState.selectedRoutes()).hasValue(ImmutableList.of(new Route.Tag("next"))); + assertThat(context.eventAuthor()).isEqualTo("custom-author"); + assertThat(context.routes()).hasValue(ImmutableList.of(new Route.Tag("next"))); + assertThat(context.output()).hasValue("done"); + } + + @Test + public void run_agentNodeOnABranch_runsTheAgentOnThatBranch() { + TestBaseAgent agent = new TestBaseAgent("helper", "An agent.", null, null, Flowable::empty); + + NodeRunner.builder() + .node(agent) + .invocationContext(invocationContext(/* branch= */ "main")) + .eventSink(recordingSink) + .build() + .run(/* nodeInput= */ null) + .blockingSubscribe(); + + assertThat(agent.getLastInvocationContext().branch()).hasValue("main"); + } + + @Test + public void run_outputProducedOnTheNodeState_isSentOnALastEvent() { + Context context = + runAtRoot( + new TestNode( + "worker", + (ctx, input) -> { + ctx.requireNodeState().produceOutput("result"); + return Flowable.empty(); + })); + + Event event = getOnlyElement(sentEvents); + assertThat(event.output()).hasValue("result"); + assertThat(event.nodeInfo().map(NodeInfo::outputFor)).hasValue(ImmutableList.of("worker@1")); + assertThat(context.requireNodeState().outputEmitted()).isTrue(); + } + + @Test + public void run_stateChangeWithoutAnEvent_isSentOnALastEvent() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> { + ctx.state().put("key", "value"); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().stateDelta()).containsExactly("key", "value"); + } + + @Test + public void run_stateChange_ridesTheNextFinalEventOnly() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> + Flowable.defer( + () -> { + ctx.state().put("key", "value"); + return Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()); + }))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).actions().stateDelta()).isEmpty(); + assertThat(sentEvents.get(1).actions().stateDelta()).containsExactly("key", "value"); + } + + @Test + public void run_failure_sendsAnErrorEventWithThePendingStateChangeOnly() { + IllegalStateException failure = new IllegalStateException("node failed"); + + Context context = + runAtRoot( + new TestNode( + "boom", + (ctx, input) -> { + ctx.requireNodeState().produceOutput("unsent"); + ctx.requireNodeState().selectRoutes(ImmutableList.of(new Route.Tag("unsent"))); + ctx.state().put("key", "value"); + return Flowable.error(failure); + })); + + Event errorEvent = getOnlyElement(sentEvents); + assertThat(errorEvent.author()).isEqualTo("boom"); + assertThat(errorEvent.nodeInfo().map(NodeInfo::path)).hasValue("boom@1"); + assertThat(errorEvent.errorCode()).hasValue(ILLEGAL_STATE); + assertThat(errorEvent.errorMessage()).hasValue("node failed"); + assertThat(errorEvent.actions().stateDelta()).containsExactly("key", "value"); + assertThat(errorEvent.output()).isEmpty(); + assertThat(errorEvent.actions().route()).isEmpty(); + assertThat(context.requireNodeState().failure()).hasValue(failure); + } + + @Test + public void run_artifactChangeWithoutAnEvent_isSentOnALastEvent() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.empty(); + })); + + assertThat(getOnlyElement(sentEvents).actions().artifactDelta()) + .containsExactly("report.txt", 1); + } + + @Test + public void run_artifactChange_ridesTheNextFinalEventOnly() { + runAtRoot( + new TestNode( + "writer", + (ctx, input) -> + Flowable.defer( + () -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.just( + newEvent().content(MESSAGE).partial(true).build(), + newEvent().content(MESSAGE).build()); + }))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).actions().artifactDelta()).isEmpty(); + assertThat(sentEvents.get(1).actions().artifactDelta()).containsExactly("report.txt", 1); + } + + @Test + public void run_failure_sendsAnErrorEventWithThePendingArtifactChange() { + runAtRoot( + new TestNode( + "boom", + (ctx, input) -> { + ctx.eventActions().artifactDelta().put("report.txt", 1); + return Flowable.error(new IllegalStateException("node failed")); + })); + + Event errorEvent = getOnlyElement(sentEvents); + assertThat(errorEvent.errorCode()).hasValue(ILLEGAL_STATE); + assertThat(errorEvent.actions().artifactDelta()).containsExactly("report.txt", 1); + } + + @Test + public void run_failureAfterEvents_sendsTheErrorEventOnceTheEventsAreProcessed() { + Node node = + new TestNode( + "boom", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).build(), newEvent().content(MESSAGE).build()) + .concatWith(Flowable.error(new IllegalStateException("node failed")))); + + TestObserver observer = + rootRunner(node, acknowledgingSink).run(/* nodeInput= */ null).test(); + int sentBeforeAcknowledgements = sentEvents.size(); + acknowledgements.get(0).onComplete(); + int sentAfterFirstAcknowledgement = sentEvents.size(); + acknowledgements.get(1).onComplete(); + acknowledgements.get(2).onComplete(); + + assertThat(sentBeforeAcknowledgements).isEqualTo(1); + assertThat(sentAfterFirstAcknowledgement).isEqualTo(2); + assertThat(sentEvents).hasSize(3); + assertThat(sentEvents.get(2).errorCode()).hasValue(ILLEGAL_STATE); + observer.assertComplete(); + } + + @Test + public void run_secondOutput_failsTheNode() { + Context context = runAtRoot(new TestNode("twice", (ctx, input) -> Flowable.just("one", "two"))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).output()).hasValue("one"); + assertThat(sentEvents.get(1).errorCode()).hasValue(ILLEGAL_STATE); + assertThat(context.requireNodeState().failure()).isPresent(); + } + + @Test + public void run_error_propagatesWithoutAnErrorEvent() { + rootRunner( + new TestNode("fatal", (ctx, input) -> Flowable.error(new AssertionError("fatal"))), + recordingSink) + .run(/* nodeInput= */ null) + .test() + .assertError(AssertionError.class); + + assertThat(sentEvents).isEmpty(); + } + + @Test + public void run_transferOnTheNodesEvent_isRecordedOnItsActions() { + Context context = + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent() + .actions(EventActions.builder().transferToAgent("specialist").build()) + .build()))); + + assertThat(context.eventActions().transferToAgent()).hasValue("specialist"); + } + + @Test + public void run_eventOfANestedNode_neitherRoutesNorTransfersTheNode() { + Context context = + runAtRoot( + new TestNode( + "outer", + (ctx, input) -> + Flowable.just( + newEvent() + .author("inner") + .actions( + EventActions.builder() + .route(ImmutableList.of(new Route.Tag("next"))) + .transferToAgent("specialist") + .build()) + .build()))); + + assertThat(context.requireNodeState().selectedRoutes()).isEmpty(); + assertThat(context.eventActions().transferToAgent()).isEmpty(); + } + + @Test + public void run_routeSelectedAfterAnEventCarriedOne_isSentOnALastEvent() { + ImmutableList second = ImmutableList.of(new Route.Tag("second")); + runAtRoot( + new TestNode( + "router", + (ctx, input) -> + Flowable.just( + newEvent() + .actions( + EventActions.builder() + .route(ImmutableList.of(new Route.Tag("first"))) + .build()) + .build()) + .concatWith( + Flowable.defer( + () -> { + ctx.requireNodeState().selectRoutes(second); + return Flowable.empty(); + })))); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(1).actions().route()).hasValue(second); + } + + @Test + public void run_eventWithItsOwnBranch_keepsItsBranch() { + Context context = + runUnderParent( + new TestNode( + "worker", + (ctx, input) -> Flowable.just(newEvent().branch("own").content(MESSAGE).build())), + /* useSubBranch= */ false, + /* overrideBranch= */ null); + + assertThat(context.invocationContext().branch()).hasValue("root"); + assertThat(getOnlyElement(sentEvents).branch()).hasValue("own"); + } + + @Test + public void run_longRunningToolCalls_becomeInterruptIds() { + Context context = + runAtRoot( + new TestNode( + "asker", + (ctx, input) -> + Flowable.just( + newEvent().longRunningToolIds(ImmutableSet.of("call-1")).build()))); + + assertThat(context.requireNodeState().interruptIds()).containsExactly("call-1"); + } + + @Test + public void run_sendsEachEventAfterThePreviousOneIsProcessed() { + Node node = + new TestNode( + "pair", + (ctx, input) -> + Flowable.just( + newEvent().content(MESSAGE).build(), newEvent().content(MESSAGE).build())); + + TestObserver observer = + rootRunner(node, acknowledgingSink).run(/* nodeInput= */ null).test(); + int sentBeforeFirstAcknowledgement = sentEvents.size(); + acknowledgements.get(0).onComplete(); + int sentBeforeSecondAcknowledgement = sentEvents.size(); + observer.assertNotComplete(); + acknowledgements.get(1).onComplete(); + + assertThat(sentBeforeFirstAcknowledgement).isEqualTo(1); + assertThat(sentBeforeSecondAcknowledgement).isEqualTo(2); + observer.assertComplete(); + } + + @Test + public void run_agentNode_keepsTheAuthorOfTheAgentsEvents() { + // Agent name differs from emitted author, so stamping with the node name would fail. + Node agent = + TestUtils.createSubAgent( + "router", Event.builder().id("e1").author("worker").content(MESSAGE).build()); + + runAtRoot(agent); + + Event event = getOnlyElement(sentEvents); + assertThat(event.author()).isEqualTo("worker"); + assertThat(event.nodeInfo().map(NodeInfo::path)).hasValue("router@1"); + } + + @Test + public void run_agentNodeWithSeveralAuthors_stampsEachEventWithItsOwnAuthor() { + Node agent = + TestUtils.createSubAgent( + "team", + Event.builder().id("e1").author("planner").content(MESSAGE).build(), + Event.builder().id("e2").author("worker").content(MESSAGE).build(), + Event.builder().id("e3").author("critic").content(MESSAGE).build()); + + TestObserver observer = + rootRunner(agent, acknowledgingSink).run(/* nodeInput= */ null).test(); + acknowledgements.get(0).onComplete(); + acknowledgements.get(1).onComplete(); + acknowledgements.get(2).onComplete(); + + observer.assertComplete(); + assertThat(sentEvents.stream().map(Event::author).collect(toImmutableList())) + .containsExactly("planner", "worker", "critic") + .inOrder(); + } + + @Test + public void run_agentNode_runsTheAgentWithItsCallbacks() { + var runAsyncImpl = TestCallback.returningEmpty(); + Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output")); + var beforeCallback = TestCallback.returningEmpty(); + var afterCallback = TestCallback.returning(afterCallbackContent); + TestBaseAgent agent = + new TestBaseAgent( + "agent", + "An agent.", + ImmutableList.of(beforeCallback.asBeforeAgentCallback()), + ImmutableList.of(afterCallback.asAfterAgentCallback()), + runAsyncImpl.asRunAsyncImplSupplier(MESSAGE)); + + runAtRoot(agent); + + assertThat(sentEvents).hasSize(2); + assertThat(sentEvents.get(0).content()).hasValue(MESSAGE); + assertThat(sentEvents.get(1).content()).hasValue(afterCallbackContent); + assertThat(beforeCallback.wasCalled()).isTrue(); + assertThat(runAsyncImpl.wasCalled()).isTrue(); + } + + @Test + public void run_agentNodeEventWithoutAnAuthor_keepsThePreviousAuthor() { + Node agent = + TestUtils.createSubAgent( + "router", + Event.builder().id("e1").author("worker").content(MESSAGE).build(), + Event.builder().id("e2").author("").content(MESSAGE).build()); + + runAtRoot(agent); + + assertThat(sentEvents.stream().map(Event::author).collect(toImmutableList())) + .containsExactly("worker", "worker") + .inOrder(); + } + + @Test + public void build_withoutANode_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .invocationContext(invocationContext(/* branch= */ null)) + .eventSink(recordingSink); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_atTheRootWithoutAnEventSink_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .node(TestNode.of("worker")) + .invocationContext(invocationContext(/* branch= */ null)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @Test + public void build_underAParentWithAnInvocationContext_throwsIllegalStateException() { + NodeRunner.Builder builder = + NodeRunner.builder() + .node(TestNode.of("worker")) + .parent(parentContext()) + .invocationContext(invocationContext(/* branch= */ null)); + + assertThrows(IllegalStateException.class, builder::build); + } + + @CanIgnoreReturnValue + private Context runAtRoot(Node node) { + return rootRunner(node, recordingSink).run(/* nodeInput= */ null).blockingGet(); + } + + private static NodeRunner rootRunner(Node node, NodeRunner.EventSink eventSink) { + return NodeRunner.builder() + .node(node) + .invocationContext(invocationContext(/* branch= */ null)) + .eventSink(eventSink) + .build(); + } + + private Context runUnderParent(Node node, boolean useSubBranch, @Nullable String overrideBranch) { + return NodeRunner.builder() + .node(node) + .parent(parentContext()) + .useSubBranch(useSubBranch) + .overrideBranch(overrideBranch) + .build() + .run(/* nodeInput= */ null) + .blockingGet(); + } + + @CanIgnoreReturnValue + private Context runWorkerUnder(Context parent) { + return NodeRunner.builder() + .node(new TestNode("worker", (ctx, input) -> Flowable.just("done"))) + .parent(parent) + .runId("2") + .build() + .run(/* nodeInput= */ null) + .blockingGet(); + } + + /** Returns the context of a node named {@code root} that runs on the {@code root} branch. */ + private Context parentContext() { + return new Context( + invocationContext("root"), + NodeExecutionState.builder() + .node(TestNode.of("root")) + .eventSink(recordingSink) + .runId("1") + .build()); + } + + private static InvocationContext invocationContext(@Nullable String branch) { + return TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .branch(branch) + .build(); + } + + private static Event.Builder newEvent() { + return Event.builder().id(Event.generateEventId()).author(""); + } + + private static Event messageAsOutput(@Nullable Object output) { + return newEvent() + .content(MESSAGE) + .output(output) + .nodeInfo(NodeInfo.builder().messageAsOutput(true).build()) + .build(); + } +} diff --git a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java index cf1f90aa2..3e2b93e08 100644 --- a/core/src/test/java/com/google/adk/workflow/WorkflowTest.java +++ b/core/src/test/java/com/google/adk/workflow/WorkflowTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertThrows; import com.google.adk.agents.CallbackContext; +import com.google.adk.testing.TestNode; import com.google.adk.testing.TestUtils; import com.google.common.collect.ImmutableList; import com.google.genai.types.Schema; @@ -35,8 +36,8 @@ public final class WorkflowTest { @Test public void build_buildsTheGraphFromTheEdges() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); ImmutableList edges = ImmutableList.of(new Edge(Start.INSTANCE, a), new Edge(a, b)); Workflow workflow = Workflow.builder().name("flow").edges(edges).build(); @@ -103,9 +104,9 @@ public void build_setsEveryProperty() { @Test public void edges_withVarargs_setsTheEdgesInOrder() { - Node a = new StubNode("a"); + Node a = TestNode.of("a"); Edge first = new Edge(Start.INSTANCE, a); - Edge second = new Edge(a, new StubNode("b")); + Edge second = new Edge(a, TestNode.of("b")); Workflow workflow = Workflow.builder().name("flow").edges(first, second).build(); @@ -114,8 +115,8 @@ public void edges_withVarargs_setsTheEdgesInOrder() { @Test public void edges_withAConsumer_setsTheEdgesItAdds() { - Node a = new StubNode("a"); - Node b = new StubNode("b"); + Node a = TestNode.of("a"); + Node b = TestNode.of("b"); Workflow workflow = Workflow.builder().name("flow").edges(g -> g.fromStart().then(a).then(b)).build(); @@ -189,7 +190,7 @@ public void build_rejectsAMissingName() { @Test public void aWorkflowCanBeANodeOfAnotherWorkflow() { Workflow inner = - Workflow.builder().name("inner").edges(new Edge(Start.INSTANCE, new StubNode("a"))).build(); + Workflow.builder().name("inner").edges(new Edge(Start.INSTANCE, TestNode.of("a"))).build(); Workflow outer = Workflow.builder().name("outer").edges(new Edge(Start.INSTANCE, inner)).build(); @@ -200,7 +201,7 @@ public void aWorkflowCanBeANodeOfAnotherWorkflow() { @Test public void runNode_failsAsNotSupportedYet() { Workflow workflow = - Workflow.builder().name("flow").edges(new Edge(Start.INSTANCE, new StubNode("a"))).build(); + Workflow.builder().name("flow").edges(new Edge(Start.INSTANCE, TestNode.of("a"))).build(); CallbackContext context = new CallbackContext( TestUtils.createInvocationContext(TestUtils.createRootAgent()),