refactor: PluginManager is used as a Plugin interface where possible

InvocationContext -> PluginManager causes some interesting circular dependencies.  Next step should probably be an interface for InvocationContext that can be used by Plugin.

PiperOrigin-RevId: 853264745

Author: google-genai-bot

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

@@ -22,7 +22,7 @@
 import com.google.adk.agents.Callbacks.AfterAgentCallback;
 import com.google.adk.agents.Callbacks.BeforeAgentCallback;
 import com.google.adk.events.Event;
-import com.google.adk.plugins.PluginManager;
+import com.google.adk.plugins.Plugin;
 import com.google.common.collect.ImmutableList;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.DoNotCall;
@@ -255,9 +255,9 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
    * @return callback functions.
    */
   private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacksToFunctions(
-      PluginManager pluginManager, List<? extends BeforeAgentCallback> callbacks) {
+      Plugin pluginManager, List<? extends BeforeAgentCallback> callbacks) {
     return Stream.concat(
-            Stream.of(ctx -> pluginManager.runBeforeAgentCallback(this, ctx)),
+            Stream.of(ctx -> pluginManager.beforeAgentCallback(this, ctx)),
             callbacks.stream()
                 .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
         .collect(toImmutableList());
@@ -270,9 +270,9 @@ private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacks
    * @return callback functions.
    */
   private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksToFunctions(
-      PluginManager pluginManager, List<? extends AfterAgentCallback> callbacks) {
+      Plugin pluginManager, List<? extends AfterAgentCallback> callbacks) {
     return Stream.concat(
-            Stream.of(ctx -> pluginManager.runAfterAgentCallback(this, ctx)),
+            Stream.of(ctx -> pluginManager.afterAgentCallback(this, ctx)),
             callbacks.stream()
                 .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
         .collect(toImmutableList());

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

@@ -21,6 +21,7 @@
 import com.google.adk.flows.llmflows.ResumabilityConfig;
 import com.google.adk.memory.BaseMemoryService;
 import com.google.adk.models.LlmCallsLimitExceededException;
+import com.google.adk.plugins.Plugin;
 import com.google.adk.plugins.PluginManager;
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
@@ -42,7 +43,7 @@ public class InvocationContext {
   private final BaseSessionService sessionService;
   private final BaseArtifactService artifactService;
   private final BaseMemoryService memoryService;
-  private final PluginManager pluginManager;
+  private final Plugin pluginManager;
   private final Optional<LiveRequestQueue> liveRequestQueue;
   private final Map<String, ActiveStreamingTool> activeStreamingTools = new ConcurrentHashMap<>();
   private final String invocationId;
@@ -80,7 +81,7 @@ public InvocationContext(
       BaseSessionService sessionService,
       BaseArtifactService artifactService,
       BaseMemoryService memoryService,
-      PluginManager pluginManager,
+      Plugin pluginManager,
       Optional<LiveRequestQueue> liveRequestQueue,
       Optional<String> branch,
       String invocationId,
@@ -235,7 +236,7 @@ public BaseMemoryService memoryService() {
   }
 
   /** Returns the plugin manager for accessing tools and plugins. */
-  public PluginManager pluginManager() {
+  public Plugin pluginManager() {
     return pluginManager;
   }
 
@@ -376,7 +377,7 @@ public static class Builder {
     private BaseSessionService sessionService;
     private BaseArtifactService artifactService;
     private BaseMemoryService memoryService;
-    private PluginManager pluginManager = new PluginManager();
+    private Plugin pluginManager = new PluginManager();
     private Optional<LiveRequestQueue> liveRequestQueue = Optional.empty();
     private Optional<String> branch = Optional.empty();
     private String invocationId = newInvocationContextId();
@@ -430,7 +431,7 @@ public Builder memoryService(BaseMemoryService memoryService) {
      * @return this builder instance for chaining.
      */
     @CanIgnoreReturnValue
-    public Builder pluginManager(PluginManager pluginManager) {
+    public Builder pluginManager(Plugin pluginManager) {
       this.pluginManager = pluginManager;
       return this;
     }

core/src/main/java/com/google/adk/flows/llmflows/BaseLlmFlow.java

@@ -200,7 +200,7 @@ private Flowable<LlmResponse> callLlm(
                                   exception ->
                                       context
                                           .pluginManager()
-                                          .runOnModelErrorCallback(
+                                          .onModelErrorCallback(
                                               new CallbackContext(
                                                   context, eventForCallbackUsage.actions()),
                                               llmRequestBuilder,
@@ -244,7 +244,7 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
     CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
 
     Maybe<LlmResponse> pluginResult =
-        context.pluginManager().runBeforeModelCallback(callbackContext, llmRequestBuilder);
+        context.pluginManager().beforeModelCallback(callbackContext, llmRequestBuilder);
 
     LlmAgent agent = (LlmAgent) context.agent();
 
@@ -280,7 +280,7 @@ private Single<LlmResponse> handleAfterModelCallback(
     CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
 
     Maybe<LlmResponse> pluginResult =
-        context.pluginManager().runAfterModelCallback(callbackContext, llmResponse);
+        context.pluginManager().afterModelCallback(callbackContext, llmResponse);
 
     LlmAgent agent = (LlmAgent) context.agent();
     Optional<List<? extends AfterModelCallback>> callbacksOpt = agent.afterModelCallback();

core/src/main/java/com/google/adk/flows/llmflows/Functions.java

@@ -390,7 +390,7 @@ private static Maybe<Event> postProcessFunctionResult(
             t ->
                 invocationContext
                     .pluginManager()
-                    .runOnToolErrorCallback(tool, functionArgs, toolContext, t)
+                    .onToolErrorCallback(tool, functionArgs, toolContext, t)
                     .map(isLive ? Optional::ofNullable : Optional::of)
                     .switchIfEmpty(Single.error(t)))
         .flatMapMaybe(
@@ -462,7 +462,7 @@ private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
       LlmAgent agent = (LlmAgent) invocationContext.agent();
 
       Maybe<Map<String, Object>> pluginResult =
-          invocationContext.pluginManager().runBeforeToolCallback(tool, functionArgs, toolContext);
+          invocationContext.pluginManager().beforeToolCallback(tool, functionArgs, toolContext);
 
       Optional<List<? extends BeforeToolCallback>> callbacksOpt = agent.beforeToolCallback();
       if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
@@ -496,7 +496,7 @@ private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
       Maybe<Map<String, Object>> pluginResult =
           invocationContext
               .pluginManager()
-              .runAfterToolCallback(tool, functionArgs, toolContext, functionResult);
+              .afterToolCallback(tool, functionArgs, toolContext, functionResult);
 
       Optional<List<? extends AfterToolCallback>> callbacksOpt = agent.afterToolCallback();
       if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {

core/src/main/java/com/google/adk/plugins/BasePlugin.java

@@ -15,19 +15,6 @@
  */
 package com.google.adk.plugins;
 
-import com.google.adk.agents.BaseAgent;
-import com.google.adk.agents.CallbackContext;
-import com.google.adk.agents.InvocationContext;
-import com.google.adk.events.Event;
-import com.google.adk.models.LlmRequest;
-import com.google.adk.models.LlmResponse;
-import com.google.adk.tools.BaseTool;
-import com.google.adk.tools.ToolContext;
-import com.google.genai.types.Content;
-import io.reactivex.rxjava3.core.Completable;
-import io.reactivex.rxjava3.core.Maybe;
-import java.util.Map;
-
 /**
  * Base class for creating plugins.
  *
@@ -40,168 +27,15 @@
  * <p>A plugin can implement one or more methods of callbacks, but should not implement the same
  * method of callback for multiple times.
  */
-public abstract class BasePlugin {
+public abstract class BasePlugin implements Plugin {
   protected final String name;
 
   public BasePlugin(String name) {
     this.name = name;
   }
 
+  @Override
   public String getName() {
     return name;
   }
-
-  /**
-   * Callback executed when a user message is received before an invocation starts.
-   *
-   * @param invocationContext The context for the entire invocation.
-   * @param userMessage The message content input by user.
-   * @return An optional Content to replace the user message. Returning Empty to proceed normally.
-   */
-  public Maybe<Content> onUserMessageCallback(
-      InvocationContext invocationContext, Content userMessage) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed before the ADK runner runs.
-   *
-   * @param invocationContext The context for the entire invocation.
-   * @return An optional Content to halt execution. Returning Empty to proceed normally.
-   */
-  public Maybe<Content> beforeRunCallback(InvocationContext invocationContext) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed after an event is yielded from runner.
-   *
-   * @param invocationContext The context for the entire invocation.
-   * @param event The event raised by the runner.
-   * @return An optional Event to modify or replace the response. Returning Empty to proceed
-   *     normally.
-   */
-  public Maybe<Event> onEventCallback(InvocationContext invocationContext, Event event) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed after an ADK runner run has completed.
-   *
-   * @param invocationContext The context for the entire invocation.
-   */
-  public Completable afterRunCallback(InvocationContext invocationContext) {
-    return Completable.complete();
-  }
-
-  /**
-   * Callback executed before an agent's primary logic is invoked.
-   *
-   * @param agent The agent that is about to run.
-   * @param callbackContext The context for the agent invocation.
-   * @return An optional Content object to bypass the agent's execution. Returning Empty to proceed
-   *     normally.
-   */
-  public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed after an agent's primary logic has completed.
-   *
-   * @param agent The agent that has just run.
-   * @param callbackContext The context for the agent invocation.
-   * @return An optional Content object to replace the agent's original result. Returning Empty to
-   *     use the original result.
-   */
-  public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed before a request is sent to the model.
-   *
-   * @param callbackContext The context for the current agent call.
-   * @param llmRequest The mutable request builder, allowing modification of the request before it
-   *     is sent to the model.
-   * @return An optional LlmResponse to trigger an early exit. Returning Empty to proceed normally.
-   */
-  public Maybe<LlmResponse> beforeModelCallback(
-      CallbackContext callbackContext, LlmRequest.Builder llmRequest) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed after a response is received from the model.
-   *
-   * @param callbackContext The context for the current agent call.
-   * @param llmResponse The response object received from the model.
-   * @return An optional LlmResponse to modify or replace the response. Returning Empty to use the
-   *     original response.
-   */
-  public Maybe<LlmResponse> afterModelCallback(
-      CallbackContext callbackContext, LlmResponse llmResponse) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed when a model call encounters an error.
-   *
-   * @param callbackContext The context for the current agent call.
-   * @param llmRequest The mutable request builder for the request that failed.
-   * @param error The exception that was raised.
-   * @return An optional LlmResponse to use instead of propagating the error. Returning Empty to
-   *     allow the original error to be raised.
-   */
-  public Maybe<LlmResponse> onModelErrorCallback(
-      CallbackContext callbackContext, LlmRequest.Builder llmRequest, Throwable error) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed before a tool is called.
-   *
-   * @param tool The tool instance that is about to be executed.
-   * @param toolArgs The dictionary of arguments to be used for invoking the tool.
-   * @param toolContext The context specific to the tool execution.
-   * @return An optional Map to stop the tool execution and return this response immediately.
-   *     Returning Empty to proceed normally.
-   */
-  public Maybe<Map<String, Object>> beforeToolCallback(
-      BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed after a tool has been called.
-   *
-   * @param tool The tool instance that has just been executed.
-   * @param toolArgs The original arguments that were passed to the tool.
-   * @param toolContext The context specific to the tool execution.
-   * @param result The dictionary returned by the tool invocation.
-   * @return An optional Map to replace the original result from the tool. Returning Empty to use
-   *     the original result.
-   */
-  public Maybe<Map<String, Object>> afterToolCallback(
-      BaseTool tool,
-      Map<String, Object> toolArgs,
-      ToolContext toolContext,
-      Map<String, Object> result) {
-    return Maybe.empty();
-  }
-
-  /**
-   * Callback executed when a tool call encounters an error.
-   *
-   * @param tool The tool instance that encountered an error.
-   * @param toolArgs The arguments that were passed to the tool.
-   * @param toolContext The context specific to the tool execution.
-   * @param error The exception that was raised during tool execution.
-   * @return An optional Map to be used as the tool response instead of propagating the error.
-   *     Returning Empty to allow the original error to be raised.
-   */
-  public Maybe<Map<String, Object>> onToolErrorCallback(
-      BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext, Throwable error) {
-    return Maybe.empty();
-  }
 }