refactor: Using a new InvocationContext.combinedPlugin() to simplify code

PiperOrigin-RevId: 855427656

Author: google-genai-bot

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

@@ -16,9 +16,6 @@
 
 package com.google.adk.agents;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
-import static java.util.Arrays.stream;
-
 import com.google.adk.Telemetry;
 import com.google.adk.agents.Callbacks.AfterAgentCallback;
 import com.google.adk.agents.Callbacks.BeforeAgentCallback;
@@ -255,10 +252,11 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
               spanContext,
               span,
               () ->
-                  callCallback(
-                          beforeCallbacksToFunctions(
-                              invocationContext.pluginManager(), callbackPlugin),
+                  processAgentCallbackResult(
+                          ctx -> invocationContext.combinedPlugin().beforeAgentCallback(this, ctx),
                           invocationContext)
+                      .map(Optional::of)
+                      .switchIfEmpty(Single.just(Optional.empty()))
                       .flatMapPublisher(
                           beforeEventOpt -> {
                             if (invocationContext.endInvocation()) {
@@ -271,11 +269,14 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
                             Flowable<Event> afterEvents =
                                 Flowable.defer(
                                     () ->
-                                        callCallback(
-                                                afterCallbacksToFunctions(
-                                                    invocationContext.pluginManager(),
-                                                    callbackPlugin),
+                                        processAgentCallbackResult(
+                                                ctx ->
+                                                    invocationContext
+                                                        .combinedPlugin()
+                                                        .afterAgentCallback(this, ctx),
                                                 invocationContext)
+                                            .map(Optional::of)
+                                            .switchIfEmpty(Single.just(Optional.empty()))
                                             .flatMapPublisher(Flowable::fromOptional));
 
                             return Flowable.concat(beforeEvents, mainEvents, afterEvents);
@@ -284,73 +285,32 @@ public Flowable<Event> runAsync(InvocationContext parentContext) {
   }
 
   /**
-   * Converts before-agent callbacks to functions.
+   * Processes the result of an agent callback, creating an {@link Event} if necessary.
    *
-   * @return callback functions.
+   * @param agentCallback The callback function.
+   * @param invocationContext The current invocation context.
+   * @return A {@link Maybe} emitting an {@link Event} if one is produced, or empty otherwise.
    */
-  private ImmutableList<Function<CallbackContext, Maybe<Content>>> beforeCallbacksToFunctions(
-      Plugin... plugins) {
-    return stream(plugins)
-        .map(
-            p ->
-                (Function<CallbackContext, Maybe<Content>>) ctx -> p.beforeAgentCallback(this, ctx))
-        .collect(toImmutableList());
-  }
-
-  /**
-   * Converts after-agent callbacks to functions.
-   *
-   * @return callback functions.
-   */
-  private ImmutableList<Function<CallbackContext, Maybe<Content>>> afterCallbacksToFunctions(
-      Plugin... plugins) {
-    return stream(plugins)
-        .map(
-            p -> (Function<CallbackContext, Maybe<Content>>) ctx -> p.afterAgentCallback(this, ctx))
-        .collect(toImmutableList());
-  }
-
-  /**
-   * Calls agent callbacks and returns the first produced event, if any.
-   *
-   * @param agentCallbacks Callback functions.
-   * @param invocationContext Current invocation context.
-   * @return single emitting first event, or empty if none.
-   */
-  private Single<Optional<Event>> callCallback(
-      List<Function<CallbackContext, Maybe<Content>>> agentCallbacks,
+  private Maybe<Event> processAgentCallbackResult(
+      Function<CallbackContext, Maybe<Content>> agentCallback,
       InvocationContext invocationContext) {
-    if (agentCallbacks == null || agentCallbacks.isEmpty()) {
-      return Single.just(Optional.empty());
-    }
-
-    CallbackContext callbackContext =
-        new CallbackContext(invocationContext, /* eventActions= */ null);
-
-    return Flowable.fromIterable(agentCallbacks)
-        .concatMap(
-            callback -> {
-              Maybe<Content> maybeContent = callback.apply(callbackContext);
-
-              return maybeContent
-                  .map(
-                      content -> {
-                        invocationContext.setEndInvocation(true);
-                        return Optional.of(
-                            Event.builder()
-                                .id(Event.generateEventId())
-                                .invocationId(invocationContext.invocationId())
-                                .author(name())
-                                .branch(invocationContext.branch())
-                                .actions(callbackContext.eventActions())
-                                .content(content)
-                                .build());
-                      })
-                  .toFlowable();
+    var callbackContext = new CallbackContext(invocationContext, /* eventActions= */ null);
+    return agentCallback
+        .apply(callbackContext)
+        .map(
+            content -> {
+              invocationContext.setEndInvocation(true);
+              return Event.builder()
+                  .id(Event.generateEventId())
+                  .invocationId(invocationContext.invocationId())
+                  .author(name())
+                  .branch(invocationContext.branch())
+                  .actions(callbackContext.eventActions())
+                  .content(content)
+                  .build();
             })
-        .firstElement()
         .switchIfEmpty(
-            Single.defer(
+            Maybe.defer(
                 () -> {
                   if (callbackContext.state().hasDelta()) {
                     Event.Builder eventBuilder =
@@ -361,9 +321,9 @@ private Single<Optional<Event>> callCallback(
                             .branch(invocationContext.branch())
                             .actions(callbackContext.eventActions());
 
-                    return Single.just(Optional.of(eventBuilder.build()));
+                    return Maybe.just(eventBuilder.build());
                   } else {
-                    return Single.just(Optional.empty());
+                    return Maybe.empty();
                   }
                 }));
   }

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

@@ -25,6 +25,7 @@
 import com.google.adk.plugins.PluginManager;
 import com.google.adk.sessions.BaseSessionService;
 import com.google.adk.sessions.Session;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import com.google.errorprone.annotations.InlineMe;
@@ -44,6 +45,7 @@ public class InvocationContext {
   private final BaseArtifactService artifactService;
   private final BaseMemoryService memoryService;
   private final Plugin pluginManager;
+  private final Plugin combinedPlugin;
   private final Optional<LiveRequestQueue> liveRequestQueue;
   private final Map<String, ActiveStreamingTool> activeStreamingTools;
   private final String invocationId;
@@ -73,6 +75,13 @@ protected InvocationContext(Builder builder) {
     this.endInvocation = builder.endInvocation;
     this.resumabilityConfig = builder.resumabilityConfig;
     this.invocationCostManager = builder.invocationCostManager;
+    this.combinedPlugin =
+        Optional.ofNullable(builder.agent)
+            .map(BaseAgent::getPlugin)
+            .map(
+                agentPlugin ->
+                    (Plugin) new PluginManager(ImmutableList.of(pluginManager, agentPlugin)))
+            .orElse(pluginManager);
   }
 
   /**
@@ -235,6 +244,14 @@ public Plugin pluginManager() {
     return pluginManager;
   }
 
+  /**
+   * Returns a {@link Plugin} that combines agent-specific plugins with framework-level plugins,
+   * allowing tools from both to be invoked.
+   */
+  public Plugin combinedPlugin() {
+    return combinedPlugin;
+  }
+
   /** Returns a map of tool call IDs to active streaming tools for the current invocation. */
   public Map<String, ActiveStreamingTool> activeStreamingTools() {
     return activeStreamingTools;

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

@@ -199,7 +199,7 @@ private Flowable<LlmResponse> callLlm(
                               .onErrorResumeNext(
                                   exception ->
                                       context
-                                          .pluginManager()
+                                          .combinedPlugin()
                                           .onModelErrorCallback(
                                               new CallbackContext(
                                                   context, eventForCallbackUsage.actions()),
@@ -243,27 +243,9 @@ private Single<Optional<LlmResponse>> handleBeforeModelCallback(
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
 
-    Maybe<LlmResponse> pluginResult =
-        context.pluginManager().beforeModelCallback(callbackContext, llmRequestBuilder);
-
-    LlmAgent agent = (LlmAgent) context.agent();
-
-    Optional<List<? extends BeforeModelCallback>> callbacksOpt = agent.beforeModelCallback();
-    if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-      return pluginResult.map(Optional::of).defaultIfEmpty(Optional.empty());
-    }
-
-    List<? extends BeforeModelCallback> callbacks = callbacksOpt.get();
-
-    Maybe<LlmResponse> callbackResult =
-        Maybe.defer(
-            () ->
-                Flowable.fromIterable(callbacks)
-                    .concatMapMaybe(callback -> callback.call(callbackContext, llmRequestBuilder))
-                    .firstElement());
-
-    return pluginResult
-        .switchIfEmpty(callbackResult)
+    return context
+        .combinedPlugin()
+        .beforeModelCallback(callbackContext, llmRequestBuilder)
         .map(Optional::of)
         .defaultIfEmpty(Optional.empty());
   }
@@ -279,24 +261,10 @@ private Single<LlmResponse> handleAfterModelCallback(
     Event callbackEvent = modelResponseEvent.toBuilder().build();
     CallbackContext callbackContext = new CallbackContext(context, callbackEvent.actions());
 
-    Maybe<LlmResponse> pluginResult =
-        context.pluginManager().afterModelCallback(callbackContext, llmResponse);
-
-    LlmAgent agent = (LlmAgent) context.agent();
-    Optional<List<? extends AfterModelCallback>> callbacksOpt = agent.afterModelCallback();
-
-    if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-      return pluginResult.defaultIfEmpty(llmResponse);
-    }
-
-    Maybe<LlmResponse> callbackResult =
-        Maybe.defer(
-            () ->
-                Flowable.fromIterable(callbacksOpt.get())
-                    .concatMapMaybe(callback -> callback.call(callbackContext, llmResponse))
-                    .firstElement());
-
-    return pluginResult.switchIfEmpty(callbackResult).defaultIfEmpty(llmResponse);
+    return context
+        .combinedPlugin()
+        .afterModelCallback(callbackContext, llmResponse)
+        .defaultIfEmpty(llmResponse);
   }
 
   /**

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

@@ -22,10 +22,7 @@
 
 import com.google.adk.Telemetry;
 import com.google.adk.agents.ActiveStreamingTool;
-import com.google.adk.agents.Callbacks.AfterToolCallback;
-import com.google.adk.agents.Callbacks.BeforeToolCallback;
 import com.google.adk.agents.InvocationContext;
-import com.google.adk.agents.LlmAgent;
 import com.google.adk.agents.RunConfig.ToolExecutionMode;
 import com.google.adk.events.Event;
 import com.google.adk.events.EventActions;
@@ -388,7 +385,7 @@ private static Maybe<Event> postProcessFunctionResult(
         .onErrorResumeNext(
             t ->
                 invocationContext
-                    .pluginManager()
+                    .combinedPlugin()
                     .onToolErrorCallback(tool, functionArgs, toolContext, t)
                     .map(isLive ? Optional::ofNullable : Optional::of)
                     .switchIfEmpty(Single.error(t)))
@@ -457,30 +454,7 @@ private static Maybe<Map<String, Object>> maybeInvokeBeforeToolCall(
       BaseTool tool,
       Map<String, Object> functionArgs,
       ToolContext toolContext) {
-    if (invocationContext.agent() instanceof LlmAgent) {
-      LlmAgent agent = (LlmAgent) invocationContext.agent();
-
-      Maybe<Map<String, Object>> pluginResult =
-          invocationContext.pluginManager().beforeToolCallback(tool, functionArgs, toolContext);
-
-      Optional<List<? extends BeforeToolCallback>> callbacksOpt = agent.beforeToolCallback();
-      if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-        return pluginResult;
-      }
-      List<? extends BeforeToolCallback> callbacks = callbacksOpt.get();
-
-      Maybe<Map<String, Object>> callbackResult =
-          Maybe.defer(
-              () ->
-                  Flowable.fromIterable(callbacks)
-                      .concatMapMaybe(
-                          callback ->
-                              callback.call(invocationContext, tool, functionArgs, toolContext))
-                      .firstElement());
-
-      return pluginResult.switchIfEmpty(callbackResult);
-    }
-    return Maybe.empty();
+    return invocationContext.combinedPlugin().beforeToolCallback(tool, functionArgs, toolContext);
   }
 
   private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
@@ -489,37 +463,9 @@ private static Maybe<Map<String, Object>> maybeInvokeAfterToolCall(
       Map<String, Object> functionArgs,
       ToolContext toolContext,
       Map<String, Object> functionResult) {
-    if (invocationContext.agent() instanceof LlmAgent) {
-      LlmAgent agent = (LlmAgent) invocationContext.agent();
-
-      Maybe<Map<String, Object>> pluginResult =
-          invocationContext
-              .pluginManager()
-              .afterToolCallback(tool, functionArgs, toolContext, functionResult);
-
-      Optional<List<? extends AfterToolCallback>> callbacksOpt = agent.afterToolCallback();
-      if (callbacksOpt.isEmpty() || callbacksOpt.get().isEmpty()) {
-        return pluginResult;
-      }
-      List<? extends AfterToolCallback> callbacks = callbacksOpt.get();
-
-      Maybe<Map<String, Object>> callbackResult =
-          Maybe.defer(
-              () ->
-                  Flowable.fromIterable(callbacks)
-                      .concatMapMaybe(
-                          callback ->
-                              callback.call(
-                                  invocationContext,
-                                  tool,
-                                  functionArgs,
-                                  toolContext,
-                                  functionResult))
-                      .firstElement());
-
-      return pluginResult.switchIfEmpty(callbackResult);
-    }
-    return Maybe.empty();
+    return invocationContext
+        .combinedPlugin()
+        .afterToolCallback(tool, functionArgs, toolContext, functionResult);
   }
 
   private static Maybe<Map<String, Object>> callTool(

core/src/main/java/com/google/adk/runner/Runner.java

@@ -514,7 +514,7 @@ public Flowable<Event> runAsync(
                                                                             updatedSession,
                                                                             session);
                                                                         return contextWithUpdatedSession
-                                                                            .pluginManager()
+                                                                            .combinedPlugin()
                                                                             .onEventCallback(
                                                                                 contextWithUpdatedSession,
                                                                                 registeredEvent)