feat: enhance AgentTool to support custom plugins for observability

Author: lucas

core/src/main/java/com/google/adk/tools/AgentTool.java

@@ -26,6 +26,7 @@
 import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
 import com.google.adk.agents.LlmAgent;
 import com.google.adk.events.Event;
+import com.google.adk.plugins.Plugin;
 import com.google.adk.runner.InMemoryRunner;
 import com.google.adk.runner.Runner;
 import com.google.adk.sessions.State;
@@ -46,6 +47,7 @@ public class AgentTool extends BaseTool {
 
   private final BaseAgent agent;
   private final boolean skipSummarization;
+  private final List<Plugin> plugins;
 
   public static BaseTool fromConfig(ToolArgsConfig args, String configAbsPath)
       throws ConfigurationException {
@@ -62,21 +64,32 @@ public static BaseTool fromConfig(ToolArgsConfig args, String configAbsPath)
     }
 
     BaseAgent agent = resolvedAgents.get(0);
-    return AgentTool.create(agent, args.getOrDefault("skipSummarization", false).booleanValue());
+    return AgentTool.create(
+        agent, args.getOrDefault("skipSummarization", false).booleanValue(), ImmutableList.of());
+  }
+
+  public static AgentTool create(
+      BaseAgent agent, boolean skipSummarization, List<? extends Plugin> plugins) {
+    return new AgentTool(agent, skipSummarization, plugins);
   }
 
   public static AgentTool create(BaseAgent agent, boolean skipSummarization) {
-    return new AgentTool(agent, skipSummarization);
+    return new AgentTool(agent, skipSummarization, ImmutableList.of());
   }
 
   public static AgentTool create(BaseAgent agent) {
-    return new AgentTool(agent, false);
+    return new AgentTool(agent, false, ImmutableList.of());
   }
 
   protected AgentTool(BaseAgent agent, boolean skipSummarization) {
+    this(agent, skipSummarization, ImmutableList.of());
+  }
+
+  protected AgentTool(BaseAgent agent, boolean skipSummarization, List<? extends Plugin> plugins) {
     super(agent.name(), agent.description());
     this.agent = agent;
     this.skipSummarization = skipSummarization;
+    this.plugins = ImmutableList.copyOf(plugins != null ? plugins : ImmutableList.of());
   }
 
   @VisibleForTesting
@@ -159,7 +172,7 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
       content = Content.fromParts(Part.fromText(input.toString()));
     }
 
-    Runner runner = new InMemoryRunner(this.agent, toolContext.agentName());
+    Runner runner = new InMemoryRunner(this.agent, toolContext.agentName(), this.plugins);
     // Session state is final, can't update to toolContext state
     // session.toBuilder().setState(toolContext.getState());
     return runner
@@ -219,3 +232,4 @@ private void updateState(Map<String, Object> stateDelta, Map<String, Object> sta
         });
   }
 }
+

core/src/test/java/com/google/adk/tools/AgentToolTest.java

@@ -664,6 +664,62 @@ public void declaration_emptySequentialAgent_fallsBackToRequest() {
                 .build());
   }
 
+  @Test
+  public void create_withPlugins_initializesCorrectly() {
+    LlmAgent testAgent =
+        createTestAgentBuilder(createTestLlm(LlmResponse.builder().build()))
+            .name("agent name")
+            .description("agent description")
+            .build();
+
+    AgentTool agentTool = AgentTool.create(testAgent, false, ImmutableList.of());
+
+    assertThat(agentTool).isNotNull();
+    assertThat(agentTool.declaration()).isPresent();
+  }
+
+  @Test
+  public void runAsync_withPlugins_usesThem() {
+    LlmAgent testAgent =
+        createTestAgentBuilder(
+                createTestLlm(
+                    LlmResponse.builder()
+                        .content(Content.fromParts(Part.fromText("Sub-agent executed")))
+                        .build()))
+            .name("sub-agent")
+            .description("sub-agent description")
+            .build();
+
+    TestPlugin testPlugin = new TestPlugin();
+
+    AgentTool agentTool = AgentTool.create(testAgent, false, ImmutableList.of(testPlugin));
+
+    ToolContext toolContext = createToolContext(testAgent);
+
+    Map<String, Object> result =
+        agentTool.runAsync(ImmutableMap.of("request", "start"), toolContext).blockingGet();
+
+    assertThat(result).containsEntry("result", "Sub-agent executed");
+
+    assertThat(testPlugin.wasCalled.get()).isTrue();
+  }
+
+  private static class TestPlugin extends com.google.adk.plugins.BasePlugin {
+    final java.util.concurrent.atomic.AtomicBoolean wasCalled =
+        new java.util.concurrent.atomic.AtomicBoolean(false);
+
+    TestPlugin() {
+      super("test-plugin");
+    }
+
+    @Override
+    public Maybe<Content> onUserMessageCallback(
+        InvocationContext invocationContext, Content userMessage) {
+      wasCalled.set(true);
+      return Maybe.empty();
+    }
+  }
+
   private ToolContext createToolContext(BaseAgent agent) {
     Session session =
         sessionService.createSession("test-app", "test-user", null, "test-session").blockingGet();