feat: Integrating Plugin with ADK

This change integrates the plugin system with ADK. PluginManager is attached to the invocation context similar to session/artifact/memory.

It includes integrations with following ADK internal callbacks:

* App callbacks: Integrated in the BaseRunner class, in run_async and run_live
* On Message callbacks: Integrated in the BaseRunner class, triggers on run_async.
* Agent callbacks: Integrated in the BaseAgent class. Leveraging the existing *callback functions
* Model callbacks: Integrating in the BaseLlmFlow.
* Tool callbacks: Integrated in Function.java, wrapped around the code for agent tool_callbacks

The plugin integrations currently do not work with Bidi-streaming (live) mode.

Sample code to use plugins:

```java
# Add plugins to Runner

Runner runner = Runner(
  agent,
  "my-app",
  artifact_service,
  session_service,
  memory_service,
  ImmutableList.of(
    MySamplePlugin(),
    LoggingPlugin()),
)

```

PiperOrigin-RevId: 808029712

Author: bjacotg

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

@@ -22,6 +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.common.collect.ImmutableList;
 import com.google.errorprone.annotations.DoNotCall;
 import com.google.genai.types.Content;
@@ -34,6 +35,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
+import java.util.stream.Stream;
 import org.jspecify.annotations.Nullable;
 
 /** Base class for all agents. */
@@ -208,11 +210,11 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
             InvocationContext invocationContext = createInvocationContext(parentContext);
 
             Flowable<Event> executionFlowable =
-                beforeAgentCallback
-                    .map(
-                        callback ->
-                            callCallback(beforeCallbacksToFunctions(callback), invocationContext))
-                    .orElse(Single.just(Optional.empty()))
+                callCallback(
+                        beforeCallbacksToFunctions(
+                            invocationContext.pluginManager(),
+                            beforeAgentCallback.orElse(ImmutableList.of())),
+                        invocationContext)
                     .flatMapPublisher(
                         beforeEventOpt -> {
                           if (invocationContext.endInvocation()) {
@@ -223,16 +225,14 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
                           Flowable<Event> mainEvents =
                               Flowable.defer(() -> runAsyncImpl(invocationContext));
                           Flowable<Event> afterEvents =
-                              afterAgentCallback
-                                  .map(
-                                      callback ->
-                                          Flowable.defer(
-                                              () ->
-                                                  callCallback(
-                                                          afterCallbacksToFunctions(callback),
-                                                          invocationContext)
-                                                      .flatMapPublisher(Flowable::fromOptional)))
-                                  .orElse(Flowable.empty());
+                              Flowable.defer(
+                                  () ->
+                                      callCallback(
+                                              afterCallbacksToFunctions(
+                                                  invocationContext.pluginManager(),
+                                                  afterAgentCallback.orElse(ImmutableList.of())),
+                                              invocationContext)
+                                          .flatMapPublisher(Flowable::fromOptional));
 
                           return Flowable.concat(beforeEvents, mainEvents, afterEvents);
                         });
@@ -248,9 +248,11 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
    * @return callback functions.
    */
   private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacksToFunctions(
-      List<BeforeAgentCallback> callbacks) {
-    return callbacks.stream()
-        .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call)
+      PluginManager pluginManager, List<BeforeAgentCallback> callbacks) {
+    return Stream.concat(
+            Stream.of(ctx -> pluginManager.runBeforeAgentCallback(this, ctx)),
+            callbacks.stream()
+                .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
         .collect(toImmutableList());
   }
 
@@ -261,9 +263,11 @@ private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacks
    * @return callback functions.
    */
   private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksToFunctions(
-      List<AfterAgentCallback> callbacks) {
-    return callbacks.stream()
-        .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call)
+      PluginManager pluginManager, List<AfterAgentCallback> callbacks) {
+    return Stream.concat(
+            Stream.of(ctx -> pluginManager.runAfterAgentCallback(this, ctx)),
+            callbacks.stream()
+                .map(callback -> (Function<CallbackContext, Maybe<Content>>) callback::call))
         .collect(toImmutableList());
   }
 

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

@@ -19,6 +19,7 @@
 import com.google.adk.artifacts.BaseArtifactService;
 import com.google.adk.memory.BaseMemoryService;
 import com.google.adk.models.LlmCallsLimitExceededException;
+import com.google.adk.plugins.PluginManager;
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
 import com.google.errorprone.annotations.InlineMe;
@@ -36,6 +37,7 @@ public class InvocationContext {
   private final BaseSessionService sessionService;
   private final BaseArtifactService artifactService;
   private final BaseMemoryService memoryService;
+  private final PluginManager pluginManager;
   private final Optional<LiveRequestQueue> liveRequestQueue;
   private final Map<String, ActiveStreamingTool> activeStreamingTools = new ConcurrentHashMap<>();
 
@@ -53,6 +55,7 @@ public InvocationContext(
       BaseSessionService sessionService,
       BaseArtifactService artifactService,
       BaseMemoryService memoryService,
+      PluginManager pluginManager,
       Optional<LiveRequestQueue> liveRequestQueue,
       Optional<String> branch,
       String invocationId,
@@ -64,6 +67,7 @@ public InvocationContext(
     this.sessionService = sessionService;
     this.artifactService = artifactService;
     this.memoryService = memoryService;
+    this.pluginManager = pluginManager;
     this.liveRequestQueue = liveRequestQueue;
     this.branch = branch;
     this.invocationId = invocationId;
@@ -74,15 +78,56 @@ public InvocationContext(
     this.endInvocation = endInvocation;
   }
 
+  /**
+   * @deprecated Use the {@link #InvocationContext} constructor with PluginManager directly instead
+   */
+  @InlineMe(
+      replacement =
+          "this(sessionService, artifactService, memoryService, new"
+              + " PluginManager(), liveRequestQueue, branch, invocationId, agent,"
+              + " session, userContent, runConfig, endInvocation)",
+      imports = "com.google.adk.plugins.PluginManager")
+  @Deprecated
+  public InvocationContext(
+      BaseSessionService sessionService,
+      BaseArtifactService artifactService,
+      BaseMemoryService memoryService,
+      Optional<LiveRequestQueue> liveRequestQueue,
+      Optional<String> branch,
+      String invocationId,
+      BaseAgent agent,
+      Session session,
+      Optional<Content> userContent,
+      RunConfig runConfig,
+      boolean endInvocation) {
+    this(
+        sessionService,
+        artifactService,
+        memoryService,
+        new PluginManager(),
+        liveRequestQueue,
+        branch,
+        invocationId,
+        agent,
+        session,
+        userContent,
+        runConfig,
+        endInvocation);
+  }
+
   /**
    * @deprecated Use the {@link #InvocationContext} constructor directly instead
    */
   @InlineMe(
       replacement =
-          "new InvocationContext(sessionService, artifactService, null, Optional.empty(),"
-              + " Optional.empty(), invocationId, agent, session, Optional.ofNullable(userContent),"
-              + " runConfig, false)",
-      imports = {"com.google.adk.agents.InvocationContext", "java.util.Optional"})
+          "new InvocationContext(sessionService, artifactService, null, new PluginManager(),"
+              + " Optional.empty(), Optional.empty(), invocationId, agent, session,"
+              + " Optional.ofNullable(userContent), runConfig, false)",
+      imports = {
+        "com.google.adk.agents.InvocationContext",
+        "com.google.adk.plugins.PluginManager",
+        "java.util.Optional"
+      })
   @Deprecated
   public static InvocationContext create(
       BaseSessionService sessionService,
@@ -96,6 +141,7 @@ public static InvocationContext create(
         sessionService,
         artifactService,
         /* memoryService= */ null,
+        new PluginManager(),
         /* liveRequestQueue= */ Optional.empty(),
         /* branch= */ Optional.empty(),
         invocationId,
@@ -111,11 +157,15 @@ public static InvocationContext create(
    */
   @InlineMe(
       replacement =
-          "new InvocationContext(sessionService, artifactService, null,"
+          "new InvocationContext(sessionService, artifactService, null, new PluginManager(),"
               + " Optional.ofNullable(liveRequestQueue), Optional.empty(),"
               + " InvocationContext.newInvocationContextId(), agent, session, Optional.empty(),"
               + " runConfig, false)",
-      imports = {"com.google.adk.agents.InvocationContext", "java.util.Optional"})
+      imports = {
+        "com.google.adk.agents.InvocationContext",
+        "com.google.adk.plugins.PluginManager",
+        "java.util.Optional"
+      })
   @Deprecated
   public static InvocationContext create(
       BaseSessionService sessionService,
@@ -128,6 +178,7 @@ public static InvocationContext create(
         sessionService,
         artifactService,
         /* memoryService= */ null,
+        new PluginManager(),
         Optional.ofNullable(liveRequestQueue),
         /* branch= */ Optional.empty(),
         InvocationContext.newInvocationContextId(),
@@ -144,6 +195,7 @@ public static InvocationContext copyOf(InvocationContext other) {
             other.sessionService,
             other.artifactService,
             other.memoryService,
+            other.pluginManager,
             other.liveRequestQueue,
             other.branch,
             other.invocationId,
@@ -168,6 +220,10 @@ public BaseMemoryService memoryService() {
     return memoryService;
   }
 
+  public PluginManager pluginManager() {
+    return pluginManager;
+  }
+
   public Map<String, ActiveStreamingTool> activeStreamingTools() {
     return activeStreamingTools;
   }
@@ -260,6 +316,7 @@ public boolean equals(Object o) {
         && Objects.equals(sessionService, that.sessionService)
         && Objects.equals(artifactService, that.artifactService)
         && Objects.equals(memoryService, that.memoryService)
+        && Objects.equals(pluginManager, that.pluginManager)
         && Objects.equals(liveRequestQueue, that.liveRequestQueue)
         && Objects.equals(activeStreamingTools, that.activeStreamingTools)
         && Objects.equals(branch, that.branch)
@@ -276,6 +333,7 @@ public int hashCode() {
         sessionService,
         artifactService,
         memoryService,
+        pluginManager,
         liveRequestQueue,
         activeStreamingTools,
         branch,

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

@@ -230,6 +230,17 @@ private Flowable<LlmResponse> callLlm(
                           return llm.generateContent(
                                   llmRequestBuilder.build(),
                                   context.runConfig().streamingMode() == StreamingMode.SSE)
+                              .onErrorResumeNext(
+                                  exception ->
+                                      context
+                                          .pluginManager()
+                                          .runOnModelErrorCallback(
+                                              new CallbackContext(
+                                                  context, eventForCallbackUsage.actions()),
+                                              llmRequest,
+                                              exception)
+                                          .switchIfEmpty(Single.error(exception))
+                                          .toFlowable())
                               .doOnNext(
                                   llmResp -> {
                                     try (Scope innerScope = llmCallSpan.makeCurrent()) {
@@ -260,29 +271,32 @@ private Flowable<LlmResponse> callLlm(
    */
   private Single<Optional<LlmResponse>> handleBeforeModelCallback(
       InvocationContext context, LlmRequest.Builder llmRequestBuilder, Event modelResponseEvent) {
+    Event callbackEvent = modelResponseEvent.toBuilder().build();
+    CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
+
+    Maybe<LlmResponse> pluginResult =
+        context.pluginManager().runBeforeModelCallback(callbackContext, llmRequestBuilder.build());
+
     LlmAgent agent = (LlmAgent) context.agent();
 
     Optional<List<BeforeModelCallback>> callbacksOpt = agent.beforeModelCallback();
     if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-      return Single.just(Optional.empty());
+      return pluginResult.map(Optional::of).defaultIfEmpty(Optional.empty());
     }
 
-    Event callbackEvent = modelResponseEvent.toBuilder().build();
     List<BeforeModelCallback> callbacks = callbacksOpt.get();
 
-    return Flowable.fromIterable(callbacks)
-        .concatMapSingle(
-            callback -> {
-              CallbackContext callbackContext =
-                  new CallbackContext(context, callbackEvent.actions());
-              return callback
-                  .call(callbackContext, llmRequestBuilder)
-                  .map(Optional::of)
-                  .defaultIfEmpty(Optional.empty());
-            })
-        .filter(Optional::isPresent)
-        .firstElement()
-        .switchIfEmpty(Single.just(Optional.empty()));
+    Maybe<LlmResponse> callbackResult =
+        Maybe.defer(
+            () ->
+                Flowable.fromIterable(callbacks)
+                    .concatMapMaybe(callback -> callback.call(callbackContext, llmRequestBuilder))
+                    .firstElement());
+
+    return pluginResult
+        .switchIfEmpty(callbackResult)
+        .map(Optional::of)
+        .defaultIfEmpty(Optional.empty());
   }
 
   /**
@@ -293,30 +307,27 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
    */
   private Single<LlmResponse> handleAfterModelCallback(
       InvocationContext context, LlmResponse llmResponse, Event modelResponseEvent) {
+    Event callbackEvent = modelResponseEvent.toBuilder().build();
+    CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
+
+    Maybe<LlmResponse> pluginResult =
+        context.pluginManager().runAfterModelCallback(callbackContext, llmResponse);
+
     LlmAgent agent = (LlmAgent) context.agent();
     Optional<List<AfterModelCallback>> callbacksOpt = agent.afterModelCallback();
 
     if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-      return Single.just(llmResponse);
+      return pluginResult.defaultIfEmpty(llmResponse);
     }
 
-    Event callbackEvent = modelResponseEvent.toBuilder().content(llmResponse.content()).build();
-    List<AfterModelCallback> callbacks = callbacksOpt.get();
-
-    return Flowable.fromIterable(callbacks)
-        .concatMapSingle(
-            callback -> {
-              CallbackContext callbackContext =
-                  new CallbackContext(context, callbackEvent.actions());
-              return callback
-                  .call(callbackContext, llmResponse)
-                  .map(Optional::of)
-                  .defaultIfEmpty(Optional.empty());
-            })
-        .filter(Optional::isPresent)
-        .firstElement()
-        .map(Optional::get)
-        .switchIfEmpty(Single.just(llmResponse));
+    Maybe<LlmResponse> callbackResult =
+        Maybe.defer(
+            () ->
+                Flowable.fromIterable(callbacksOpt.get())
+                    .concatMapMaybe(callback -> callback.call(callbackContext, llmResponse))
+                    .firstElement());
+
+    return pluginResult.switchIfEmpty(callbackResult).defaultIfEmpty(llmResponse);
   }
 
   /**