feat: make BaseAgent extend BaseNode so every agent is a workflow node

PiperOrigin-RevId: 997839818

Author: damianmomotgoogle

core/src/main/java/com/google/adk/agents/BaseAgent.java

@@ -18,18 +18,21 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Strings.isNullOrEmpty;
+import static com.google.common.base.Strings.nullToEmpty;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.lang.String.format;
 
 import com.google.adk.agents.Callbacks.AfterAgentCallback;
 import com.google.adk.agents.Callbacks.BeforeAgentCallback;
+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.telemetry.Instrumentation;
 import com.google.adk.telemetry.Instrumentation.AgentInvocation;
 import com.google.adk.telemetry.Tracing;
 import com.google.adk.utils.AgentEnums.AgentOrigin;
+import com.google.adk.workflow.BaseNode;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
@@ -49,21 +52,12 @@
 import org.jspecify.annotations.Nullable;
 
 /** Base class for all agents. */
-public abstract class BaseAgent {
+public abstract class BaseAgent extends BaseNode {
 
   // Pattern for valid agent names.
   private static final String IDENTIFIER_REGEX = "^_?[a-zA-Z0-9]*([. _-][a-zA-Z0-9]+)*$";
   private static final Pattern IDENTIFIER_PATTERN = Pattern.compile(IDENTIFIER_REGEX);
 
-  /** The agent's name. Must be a unique identifier within the agent tree. */
-  private final String name;
-
-  /**
-   * One line description about the agent's capability. The system can use this for decision-making
-   * when delegating control to different agents.
-   */
-  private final String description;
-
   /**
    * The parent agent in the agent tree. Note that one agent cannot be added to two different
    * parents' sub-agents lists.
@@ -79,7 +73,7 @@ public abstract class BaseAgent {
    * Creates a new BaseAgent.
    *
    * @param name Unique agent name. Cannot be "user" (reserved).
-   * @param description Agent purpose.
+   * @param description Agent purpose; null is treated as an empty string.
    * @param subAgents Agents managed by this agent.
    * @param beforeAgentCallback Callbacks before agent execution. Invoked in order until one doesn't
    *     return null.
@@ -88,16 +82,14 @@ public abstract class BaseAgent {
    */
   public BaseAgent(
       String name,
-      String description,
+      @Nullable String description,
       @Nullable List<? extends BaseAgent> subAgents,
       @Nullable List<? extends BeforeAgentCallback> beforeAgentCallback,
       @Nullable List<? extends AfterAgentCallback> afterAgentCallback) {
-    validateAgentName(name);
-    this.name = name;
-    this.description = description;
+    super(validateAgentName(name), nullToEmpty(description));
     this.parentAgent = null;
     this.subAgents = (subAgents != null) ? ImmutableList.copyOf(subAgents) : ImmutableList.of();
-    validateSubAgents(this.name, this.subAgents);
+    validateSubAgents(name, this.subAgents);
     this.beforeAgentCallback =
         (beforeAgentCallback != null)
             ? ImmutableList.copyOf(beforeAgentCallback)
@@ -128,10 +120,12 @@ public Completable close() {
    * Validates the agent name.
    *
    * @param name The agent name to validate.
+   * @return the validated name.
    * @throws IllegalArgumentException if the agent name is null, empty, or does not match the
    *     identifier pattern.
    */
-  private static void validateAgentName(String name) {
+  @CanIgnoreReturnValue
+  private static String validateAgentName(String name) {
     if (isNullOrEmpty(name)) {
       throw new IllegalArgumentException("Agent name cannot be null or empty.");
     }
@@ -141,6 +135,7 @@ private static void validateAgentName(String name) {
     }
     checkArgument(
         !name.equals(Role.USER), "Agent name cannot be 'user'; reserved for end-user input.");
+    return name;
   }
 
   /**
@@ -173,24 +168,6 @@ private static void validateSubAgents(
     }
   }
 
-  /**
-   * Gets the agent's unique name.
-   *
-   * @return the unique name of the agent.
-   */
-  public final String name() {
-    return name;
-  }
-
-  /**
-   * Gets the one-line description of the agent's capability.
-   *
-   * @return the description of the agent.
-   */
-  public final String description() {
-    return description;
-  }
-
   /**
    * Retrieves the parent agent in the agent tree.
    *
@@ -312,7 +289,7 @@ private InvocationContext createInvocationContext(InvocationContext parentContex
    * @return stream of agent-generated events.
    */
   public Flowable<Event> runAsync(InvocationContext parentContext) {
-    return run(parentContext, this::runAsyncImpl);
+    return runWithCallbacks(parentContext, this::runAsyncImpl);
   }
 
   /**
@@ -322,7 +299,7 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
    * @param runImplementation The agent-specific logic to run.
    * @return stream of agent-generated events.
    */
-  private Flowable<Event> run(
+  private Flowable<Event> runWithCallbacks(
       InvocationContext parentContext,
       Function<InvocationContext, Flowable<Event>> runImplementation) {
     return Flowable.using(
@@ -468,7 +445,7 @@ private Maybe<Event> callCallback(
    * @return stream of agent-generated events.
    */
   public Flowable<Event> runLive(InvocationContext parentContext) {
-    return run(parentContext, this::runLiveImpl);
+    return runWithCallbacks(parentContext, this::runLiveImpl);
   }
 
   /**
@@ -510,6 +487,17 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) {
         .build();
   }
 
+  /**
+   * Runs this agent as a workflow node through its {@link #runAsync} lifecycle, including its
+   * callbacks. The node input is not used.
+   */
+  @Experimental
+  @Override
+  public Flowable<Event> runNode(
+      com.google.adk.agents.Context context, @Nullable Object nodeInput) {
+    return runAsync(context.invocationContext());
+  }
+
   /**
    * Agent-specific asynchronous logic.
    *

core/src/main/java/com/google/adk/agents/LlmAgent.java

@@ -1029,10 +1029,12 @@ public List<? extends OnToolErrorCallback> canonicalOnToolErrorCallbacks() {
     return onToolErrorCallback;
   }
 
+  @Override
   public Optional<Schema> inputSchema() {
     return inputSchema;
   }
 
+  @Override
   public Optional<Schema> outputSchema() {
     return outputSchema;
   }

core/src/test/java/com/google/adk/agents/BaseAgentTest.java

@@ -26,6 +26,7 @@
 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;
@@ -112,6 +113,13 @@ public void constructor_setsNameAndDescription() {
     assertThat(agent.description()).isEqualTo(description);
   }
 
+  @Test
+  public void constructor_nullDescription_setsEmptyDescription() {
+    TestBaseAgent agent = new TestBaseAgent("testName", null, null, ImmutableList.of(), null, null);
+
+    assertThat(agent.description()).isEmpty();
+  }
+
   @Test
   public void findAgent_returnsCorrectAgent() {
     TestBaseAgent subSubAgent =
@@ -399,6 +407,34 @@ public void runAsync_firstAfterCallbackReturnsContent_skipsSecondAfterCallback()
     assertThat(afterCallback2.wasCalled()).isFalse();
   }
 
+  @Test
+  public void runNode_runsAgentWithItsCallbacks() {
+    var runAsyncImpl = TestCallback.<Void>returningEmpty();
+    Content runAsyncImplContent = Content.fromParts(Part.fromText("main_output"));
+    Content afterCallbackContent = Content.fromParts(Part.fromText("after_callback_output"));
+    var beforeCallback = TestCallback.<Content>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<Event> 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 =
@@ -637,6 +673,20 @@ public void runLive_firstAfterCallbackReturnsContent_skipsSecondAfterCallback()
     assertThat(afterCallback2.wasCalled()).isFalse();
   }
 
+  @Test
+  public void constructor_nullName_throwsIllegalArgumentException() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> new TestBaseAgent(null, "description", null, null, null));
+  }
+
+  @Test
+  public void constructor_emptyName_throwsIllegalArgumentException() {
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> new TestBaseAgent("", "description", null, null, null));
+  }
+
   @Test
   public void constructor_invalidName_throwsIllegalArgumentException() {
     assertThrows(