From 1e468e7aa38f592c9d63983616b8a5e2c02ce2fe Mon Sep 17 00:00:00 2001 From: Google Team Member Date: Mon, 26 Jan 2026 10:50:39 -0800 Subject: [PATCH] refactor: BaseAgent: Unwrap List from Optional, improve test coverage PiperOrigin-RevId: 861263759 --- .../java/com/google/adk/agents/BaseAgent.java | 22 +++++----- .../com/google/adk/agents/BaseAgentTest.java | 20 --------- .../adk/agents/ConfigAgentUtilsTest.java | 25 +++++------ .../com/google/adk/testing/TestCallback.java | 41 ++++++------------- 4 files changed, 34 insertions(+), 74 deletions(-) diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index 7a9f14b6b..948d5ebac 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -59,8 +59,8 @@ public abstract class BaseAgent { private final List subAgents; - private final List beforeAgentCallback; - private final List afterAgentCallback; + private final Optional> beforeAgentCallback; + private final Optional> afterAgentCallback; /** * Creates a new BaseAgent. @@ -83,9 +83,8 @@ public BaseAgent( this.description = description; this.parentAgent = null; this.subAgents = subAgents != null ? subAgents : ImmutableList.of(); - this.beforeAgentCallback = - beforeAgentCallback != null ? beforeAgentCallback : ImmutableList.of(); - this.afterAgentCallback = afterAgentCallback != null ? afterAgentCallback : ImmutableList.of(); + this.beforeAgentCallback = Optional.ofNullable(beforeAgentCallback); + this.afterAgentCallback = Optional.ofNullable(afterAgentCallback); // Establish parent relationships for all sub-agents if needed. for (BaseAgent subAgent : this.subAgents) { @@ -172,11 +171,11 @@ public List subAgents() { return subAgents; } - public List beforeAgentCallback() { + public Optional> beforeAgentCallback() { return beforeAgentCallback; } - public List afterAgentCallback() { + public Optional> afterAgentCallback() { return afterAgentCallback; } @@ -186,7 +185,7 @@ public List afterAgentCallback() { *

This method is only for use by Agent Development Kit. */ public List canonicalBeforeAgentCallbacks() { - return beforeAgentCallback; + return beforeAgentCallback.orElse(ImmutableList.of()); } /** @@ -195,7 +194,7 @@ public List canonicalBeforeAgentCallbacks() { *

This method is only for use by Agent Development Kit. */ public List canonicalAfterAgentCallbacks() { - return afterAgentCallback; + return afterAgentCallback.orElse(ImmutableList.of()); } /** @@ -240,7 +239,8 @@ public Flowable runAsync(InvocationContext parentContext) { () -> callCallback( beforeCallbacksToFunctions( - invocationContext.pluginManager(), beforeAgentCallback), + invocationContext.pluginManager(), + beforeAgentCallback.orElse(ImmutableList.of())), invocationContext) .flatMapPublisher( beforeEventOpt -> { @@ -257,7 +257,7 @@ public Flowable runAsync(InvocationContext parentContext) { callCallback( afterCallbacksToFunctions( invocationContext.pluginManager(), - afterAgentCallback), + afterAgentCallback.orElse(ImmutableList.of())), invocationContext) .flatMapPublisher(Flowable::fromOptional)); diff --git a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java index 469feb6dc..345436826 100644 --- a/core/src/test/java/com/google/adk/agents/BaseAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/BaseAgentTest.java @@ -316,24 +316,4 @@ public void canonicalCallbacks_returnsListWhenPresent() { assertThat(agent.canonicalBeforeAgentCallbacks()).containsExactly(bc); assertThat(agent.canonicalAfterAgentCallbacks()).containsExactly(ac); } - - @Test - public void runLive_invokesRunLiveImpl() { - var runLiveCallback = TestCallback.returningEmpty(); - Content runLiveImplContent = Content.fromParts(Part.fromText("live_output")); - TestBaseAgent agent = - new TestBaseAgent( - TEST_AGENT_NAME, - TEST_AGENT_DESCRIPTION, - /* beforeAgentCallbacks= */ ImmutableList.of(), - /* afterAgentCallbacks= */ ImmutableList.of(), - runLiveCallback.asRunLiveImplSupplier(runLiveImplContent)); - InvocationContext invocationContext = TestUtils.createInvocationContext(agent); - - List results = agent.runLive(invocationContext).toList().blockingGet(); - - assertThat(results).hasSize(1); - assertThat(results.get(0).content()).hasValue(runLiveImplContent); - assertThat(runLiveCallback.wasCalled()).isTrue(); - } } diff --git a/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java b/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java index 9cd535aca..4825efca1 100644 --- a/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java +++ b/core/src/test/java/com/google/adk/agents/ConfigAgentUtilsTest.java @@ -1161,25 +1161,20 @@ public void fromConfig_withConfiguredCallbacks_resolvesCallbacks() String pfx = "test.callbacks."; registry.register( - pfx + "before_agent_1", (Callbacks.BeforeAgentCallback) (unusedCtx) -> Maybe.empty()); + pfx + "before_agent_1", (Callbacks.BeforeAgentCallback) (ctx) -> Maybe.empty()); registry.register( - pfx + "before_agent_2", (Callbacks.BeforeAgentCallback) (unusedCtx) -> Maybe.empty()); + pfx + "before_agent_2", (Callbacks.BeforeAgentCallback) (ctx) -> Maybe.empty()); + registry.register(pfx + "after_agent_1", (Callbacks.AfterAgentCallback) (ctx) -> Maybe.empty()); registry.register( - pfx + "after_agent_1", (Callbacks.AfterAgentCallback) (unusedCtx) -> Maybe.empty()); + pfx + "before_model_1", (Callbacks.BeforeModelCallback) (ctx, req) -> Maybe.empty()); registry.register( - pfx + "before_model_1", - (Callbacks.BeforeModelCallback) (unusedCtx, unusedReq) -> Maybe.empty()); - registry.register( - pfx + "after_model_1", - (Callbacks.AfterModelCallback) (unusedCtx, unusedResp) -> Maybe.empty()); + pfx + "after_model_1", (Callbacks.AfterModelCallback) (ctx, resp) -> Maybe.empty()); registry.register( pfx + "before_tool_1", - (Callbacks.BeforeToolCallback) - (unusedInv, unusedTool, unusedArgs, unusedToolCtx) -> Maybe.empty()); + (Callbacks.BeforeToolCallback) (inv, tool, args, toolCtx) -> Maybe.empty()); registry.register( pfx + "after_tool_1", - (Callbacks.AfterToolCallback) - (unusedInv, unusedTool, unusedArgs, unusedToolCtx, unusedResp) -> Maybe.empty()); + (Callbacks.AfterToolCallback) (inv, tool, args, toolCtx, resp) -> Maybe.empty()); File configFile = tempFolder.newFile("with_callbacks.yaml"); Files.writeString( @@ -1209,8 +1204,10 @@ public void fromConfig_withConfiguredCallbacks_resolvesCallbacks() assertThat(agent).isInstanceOf(LlmAgent.class); LlmAgent llm = (LlmAgent) agent; - assertThat(agent.beforeAgentCallback()).hasSize(2); - assertThat(agent.afterAgentCallback()).hasSize(1); + assertThat(agent.beforeAgentCallback()).isPresent(); + assertThat(agent.beforeAgentCallback().get()).hasSize(2); + assertThat(agent.afterAgentCallback()).isPresent(); + assertThat(agent.afterAgentCallback().get()).hasSize(1); assertThat(llm.beforeModelCallback()).isPresent(); assertThat(llm.beforeModelCallback().get()).hasSize(1); diff --git a/core/src/test/java/com/google/adk/testing/TestCallback.java b/core/src/test/java/com/google/adk/testing/TestCallback.java index 434d85e6f..04f83ed9b 100644 --- a/core/src/test/java/com/google/adk/testing/TestCallback.java +++ b/core/src/test/java/com/google/adk/testing/TestCallback.java @@ -102,80 +102,63 @@ public Supplier> asRunAsyncImplSupplier(String contentText) { return asRunAsyncImplSupplier(Content.fromParts(Part.fromText(contentText))); } - /** - * Returns a {@link Supplier} that marks this callback as called and returns a {@link Flowable} - * with an event containing the given content. - */ - public Supplier> asRunLiveImplSupplier(Content content) { - return () -> - Flowable.defer( - () -> { - markAsCalled(); - return Flowable.just(Event.builder().content(content).build()); - }); - } - @SuppressWarnings("unchecked") // This cast is safe if T is Content. public BeforeAgentCallback asBeforeAgentCallback() { - return (unusedCtx) -> (Maybe) callMaybe(); + return ctx -> (Maybe) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is Content. public BeforeAgentCallbackSync asBeforeAgentCallbackSync() { - return (unusedCtx) -> (Optional) callOptional(); + return ctx -> (Optional) callOptional(); } @SuppressWarnings("unchecked") // This cast is safe if T is Content. public AfterAgentCallback asAfterAgentCallback() { - return (unusedCtx) -> (Maybe) callMaybe(); + return ctx -> (Maybe) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is Content. public AfterAgentCallbackSync asAfterAgentCallbackSync() { - return (unusedCtx) -> (Optional) callOptional(); + return ctx -> (Optional) callOptional(); } @SuppressWarnings("unchecked") // This cast is safe if T is LlmResponse. public BeforeModelCallback asBeforeModelCallback() { - return (unusedCtx, unusedReq) -> (Maybe) callMaybe(); + return (ctx, req) -> (Maybe) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is LlmResponse. public BeforeModelCallbackSync asBeforeModelCallbackSync() { - return (unusedCtx, unusedReq) -> (Optional) callOptional(); + return (ctx, req) -> (Optional) callOptional(); } @SuppressWarnings("unchecked") // This cast is safe if T is LlmResponse. public AfterModelCallback asAfterModelCallback() { - return (unusedCtx, unusedRes) -> (Maybe) callMaybe(); + return (ctx, res) -> (Maybe) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is LlmResponse. public AfterModelCallbackSync asAfterModelCallbackSync() { - return (unusedCtx, unusedRes) -> (Optional) callOptional(); + return (ctx, res) -> (Optional) callOptional(); } @SuppressWarnings("unchecked") // This cast is safe if T is Map. public BeforeToolCallback asBeforeToolCallback() { - return (unusedCtx, unusedTool, unusedToolArgs, unusedToolCtx) -> - (Maybe>) callMaybe(); + return (invCtx, tool, toolArgs, toolCtx) -> (Maybe>) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is Map. public BeforeToolCallbackSync asBeforeToolCallbackSync() { - return (unusedCtx, unusedTool, unusedToolArgs, unusedToolCtx) -> - (Optional>) callOptional(); + return (invCtx, tool, toolArgs, toolCtx) -> (Optional>) callOptional(); } @SuppressWarnings("unchecked") // This cast is safe if T is Map. public AfterToolCallback asAfterToolCallback() { - return (unusedCtx, unusedTool, unusedToolArgs, unusedToolCtx, unusedRes) -> - (Maybe>) callMaybe(); + return (invCtx, tool, toolArgs, toolCtx, res) -> (Maybe>) callMaybe(); } @SuppressWarnings("unchecked") // This cast is safe if T is Map. public AfterToolCallbackSync asAfterToolCallbackSync() { - return (unusedCtx, unusedTool, unusedToolArgs, unusedToolCtx, unusedRes) -> - (Optional>) callOptional(); + return (invCtx, tool, toolArgs, toolCtx, res) -> (Optional>) callOptional(); } }