feat: run a single workflow node

PiperOrigin-RevId: 1002809551

Author: damianmomotgoogle

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<? extends AfterAgentCallback> 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<Event> 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<Event> runWithCallbacks(
       InvocationContext parentContext,
+      boolean keepBranch,
       Function<InvocationContext, Flowable<Event>> 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<Event> callCallback(
    * @return stream of agent-generated events.
    */
   public Flowable<Event> 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<Event> 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());
+              }
+            });
   }
 
   /**

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.
  *
  * <p>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<String> functionCallId = Optional.empty();
   private Optional<ToolConfirmation> 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<SearchMemoryResponse> 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<Object> 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<ImmutableList<Route>> 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<? extends Route> routes) {
+    requireNodeState().selectRoutes(routes);
+  }
 }