refactor: simplify and deduplicate logic responsible for agent loading

PiperOrigin-RevId: 834325508

Author: google-genai-bot

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

@@ -24,6 +24,7 @@
 import com.google.adk.events.Event;
 import com.google.adk.plugins.PluginManager;
 import com.google.common.collect.ImmutableList;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
 import com.google.genai.types.Content;
 import io.opentelemetry.api.trace.Span;
@@ -389,4 +390,73 @@ public static BaseAgent fromConfig(BaseAgentConfig config, String configAbsPath)
     throw new UnsupportedOperationException(
         "BaseAgent is abstract. Override fromConfig in concrete subclasses.");
   }
+
+  /**
+   * Base Builder for all agents.
+   *
+   * @param <B> The concrete builder type.
+   */
+  public abstract static class Builder<B extends Builder<B>> {
+    protected String name;
+    protected String description;
+    protected ImmutableList<BaseAgent> subAgents;
+    protected ImmutableList<BeforeAgentCallback> beforeAgentCallback;
+    protected ImmutableList<AfterAgentCallback> afterAgentCallback;
+
+    /** This is a safe cast to the concrete builder type. */
+    @SuppressWarnings("unchecked")
+    protected B self() {
+      return (B) this;
+    }
+
+    @CanIgnoreReturnValue
+    public B name(String name) {
+      this.name = name;
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B description(String description) {
+      this.description = description;
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B subAgents(List<? extends BaseAgent> subAgents) {
+      this.subAgents = ImmutableList.copyOf(subAgents);
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B subAgents(BaseAgent... subAgents) {
+      this.subAgents = ImmutableList.copyOf(subAgents);
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B beforeAgentCallback(BeforeAgentCallback beforeAgentCallback) {
+      this.beforeAgentCallback = ImmutableList.of(beforeAgentCallback);
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B beforeAgentCallback(List<Callbacks.BeforeAgentCallbackBase> beforeAgentCallback) {
+      this.beforeAgentCallback = CallbackUtil.getBeforeAgentCallbacks(beforeAgentCallback);
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B afterAgentCallback(AfterAgentCallback afterAgentCallback) {
+      this.afterAgentCallback = ImmutableList.of(afterAgentCallback);
+      return self();
+    }
+
+    @CanIgnoreReturnValue
+    public B afterAgentCallback(List<Callbacks.AfterAgentCallbackBase> afterAgentCallback) {
+      this.afterAgentCallback = CallbackUtil.getAfterAgentCallbacks(afterAgentCallback);
+      return self();
+    }
+
+    public abstract BaseAgent build();
+  }
 }

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

@@ -29,6 +29,29 @@ public class BaseAgentConfig {
   private String agentClass;
   private List<AgentRefConfig> subAgents;
 
+  // Callback configuration (names resolved via ComponentRegistry)
+  private List<CallbackRef> beforeAgentCallbacks;
+  private List<CallbackRef> afterAgentCallbacks;
+
+  /** Reference to a callback stored in the ComponentRegistry. */
+  public static class CallbackRef {
+    private String name;
+
+    public CallbackRef() {}
+
+    public CallbackRef(String name) {
+      this.name = name;
+    }
+
+    public String name() {
+      return name;
+    }
+
+    public void setName(String name) {
+      this.name = name;
+    }
+  }
+
   /**
    * Configuration for referencing other agents (subagents). Supports both config-based references
    * (YAML files) and programmatic references (via code registry).
@@ -67,6 +90,10 @@ public void setCode(String code) {
 
   public BaseAgentConfig() {}
 
+  public BaseAgentConfig(String agentClass) {
+    this.agentClass = agentClass;
+  }
+
   /**
    * Constructor with basic fields.
    *
@@ -96,19 +123,35 @@ public void setDescription(String description) {
     this.description = description;
   }
 
-  public String agentClass() {
-    return agentClass;
-  }
-
   public void setAgentClass(String agentClass) {
     this.agentClass = agentClass;
   }
 
+  public String agentClass() {
+    return agentClass;
+  }
+
   public List<AgentRefConfig> subAgents() {
     return subAgents;
   }
 
   public void setSubAgents(List<AgentRefConfig> subAgents) {
     this.subAgents = subAgents;
   }
+
+  public List<CallbackRef> beforeAgentCallbacks() {
+    return beforeAgentCallbacks;
+  }
+
+  public void setBeforeAgentCallbacks(List<CallbackRef> beforeAgentCallbacks) {
+    this.beforeAgentCallbacks = beforeAgentCallbacks;
+  }
+
+  public List<CallbackRef> afterAgentCallbacks() {
+    return afterAgentCallbacks;
+  }
+
+  public void setAfterAgentCallbacks(List<CallbackRef> afterAgentCallbacks) {
+    this.afterAgentCallbacks = afterAgentCallbacks;
+  }
 }

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

@@ -16,6 +16,8 @@
 
 package com.google.adk.agents;
 
+import static com.google.common.base.Strings.nullToEmpty;
+
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.MapperFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -31,6 +33,8 @@
 import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Consumer;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -45,6 +49,86 @@ public final class ConfigAgentUtils {
 
   private ConfigAgentUtils() {}
 
+  /**
+   * Configures the common properties of an agent builder from the configuration.
+   *
+   * @param builder The agent builder.
+   * @param config The agent configuration.
+   * @param configAbsPath The absolute path to the config file (for resolving relative paths).
+   * @throws ConfigurationException if the configuration is invalid.
+   */
+  public static void resolveAndSetCommonAgentFields(
+      BaseAgent.Builder<?> builder, BaseAgentConfig config, String configAbsPath)
+      throws ConfigurationException {
+    if (config.name() == null || config.name().trim().isEmpty()) {
+      throw new ConfigurationException("Agent name is required");
+    }
+    builder.name(config.name());
+    builder.description(nullToEmpty(config.description()));
+
+    if (config.subAgents() != null && !config.subAgents().isEmpty()) {
+      builder.subAgents(resolveSubAgents(config.subAgents(), configAbsPath));
+    }
+
+    setBaseAgentCallbacks(config, builder::beforeAgentCallback, builder::afterAgentCallback);
+  }
+
+  /**
+   * Resolves and sets callbacks from configuration.
+   *
+   * @param refs The list of callback references from config.
+   * @param callbackBaseClass The base class of the callback.
+   * @param callbackTypeName The name of the callback type for error messages.
+   * @param builderSetter The setter method on the builder to apply the resolved callbacks.
+   * @param <T> The type of the callback.
+   * @throws ConfigurationException if a callback cannot be resolved.
+   */
+  public static <T> void resolveAndSetCallback(
+      @Nullable List<BaseAgentConfig.CallbackRef> refs,
+      Class<T> callbackBaseClass,
+      String callbackTypeName,
+      Consumer<ImmutableList<T>> builderSetter)
+      throws ConfigurationException {
+    if (refs != null) {
+      ImmutableList.Builder<T> list = ImmutableList.builder();
+      for (BaseAgentConfig.CallbackRef ref : refs) {
+        list.add(
+            ComponentRegistry.getInstance()
+                .get(ref.name(), callbackBaseClass)
+                .orElseThrow(
+                    () ->
+                        new ConfigurationException(
+                            "Invalid " + callbackTypeName + ": " + ref.name())));
+      }
+      builderSetter.accept(list.build());
+    }
+  }
+
+  /**
+   * Sets the common agent callbacks (before/after agent) from the config to the builder setters.
+   *
+   * @param config The agent configuration.
+   * @param beforeSetter The setter for before-agent callbacks.
+   * @param afterSetter The setter for after-agent callbacks.
+   * @throws ConfigurationException if a callback cannot be resolved.
+   */
+  public static void setBaseAgentCallbacks(
+      BaseAgentConfig config,
+      Consumer<ImmutableList<Callbacks.BeforeAgentCallbackBase>> beforeSetter,
+      Consumer<ImmutableList<Callbacks.AfterAgentCallbackBase>> afterSetter)
+      throws ConfigurationException {
+    resolveAndSetCallback(
+        config.beforeAgentCallbacks(),
+        Callbacks.BeforeAgentCallbackBase.class,
+        "before_agent_callback",
+        beforeSetter);
+    resolveAndSetCallback(
+        config.afterAgentCallbacks(),
+        Callbacks.AfterAgentCallbackBase.class,
+        "after_agent_callback",
+        afterSetter);
+  }
+
   /**
    * Load agent from a YAML config file path.
    *
@@ -53,7 +137,6 @@ private ConfigAgentUtils() {}
    * @throws ConfigurationException if loading fails
    */
   public static BaseAgent fromConfig(String configPath) throws ConfigurationException {
-
     File configFile = new File(configPath);
     if (!configFile.exists()) {
       logger.error("Config file not found: {}", configPath);