diff --git a/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPlugin.java b/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPlugin.java new file mode 100644 index 000000000..a676ec0a7 --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPlugin.java @@ -0,0 +1,215 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.adk.plugins.BasePlugin; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.base.Strings; +import com.google.common.collect.ImmutableMap; +import io.reactivex.rxjava3.core.Maybe; +import java.util.Map; + +/** + * Provides self-healing error recovery for tool failures. + * + *

This plugin intercepts tool failures, hands the model structured guidance for reflection and + * correction, and lets it retry up to a configurable limit. Failure counts are tracked per tool + * within a scope, so a success with one tool resets that tool's counter without forgiving + * another's. + * + *

Port of adk-python's {@code ReflectAndRetryToolPlugin} ({@code + * plugins/reflect_retry_tool_plugin.py}). + * + *

Example: + * + *

{@code
+ * Runner runner =
+ *     new Runner(
+ *         agent,
+ *         APP_NAME,
+ *         artifactService,
+ *         sessionService,
+ *         ImmutableList.of(new ReflectAndRetryToolPlugin(3)));
+ * }
+ * + *

{@link #scopeKey} and {@link #extractErrorFromResult} are {@code protected} because adk-python + * documents both as overridable. + */ +public class ReflectAndRetryToolPlugin extends BasePlugin { + + private static final String DEFAULT_NAME = "reflect_retry_tool_plugin"; + private static final int DEFAULT_MAX_RETRIES = 3; + private static final String GLOBAL_SCOPE_KEY = "__global_reflect_and_retry_scope__"; + private static final String NEGATIVE_RETRIES = "maxRetries must be non-negative, but was %s"; + + /** The observed failure count when retrying is disabled: the one failure being reported. */ + private static final int FIRST_FAILURE = 1; + + private final int maxRetries; + private final boolean throwExceptionIfRetryExceeded; + private final TrackingScope trackingScope; + private final ToolFailureTracker failures = new ToolFailureTracker(); + + /** Three retries, throwing when exceeded, tracked per invocation. */ + public ReflectAndRetryToolPlugin() { + this(DEFAULT_NAME, DEFAULT_MAX_RETRIES, true, TrackingScope.INVOCATION); + } + + /** As above, with a custom retry limit. */ + public ReflectAndRetryToolPlugin(int maxRetries) { + this(DEFAULT_NAME, maxRetries, true, TrackingScope.INVOCATION); + } + + /** + * @param name plugin instance identifier + * @param maxRetries maximum consecutive failures before giving up; {@code 0} disables retrying + * @param throwExceptionIfRetryExceeded whether to propagate the final error once the limit is + * reached, rather than returning guidance + * @param trackingScope lifetime of the failure counters + * @throws IllegalArgumentException if {@code maxRetries} is negative + * @throws NullPointerException if {@code trackingScope} is null + */ + public ReflectAndRetryToolPlugin( + String name, + int maxRetries, + boolean throwExceptionIfRetryExceeded, + TrackingScope trackingScope) { + super(name); + checkArgument(maxRetries >= 0, NEGATIVE_RETRIES, maxRetries); + this.maxRetries = maxRetries; + this.throwExceptionIfRetryExceeded = throwExceptionIfRetryExceeded; + this.trackingScope = checkNotNull(trackingScope, "trackingScope cannot be null"); + } + + /** + * Resets the tool's failure count on success, or routes an error carried inside an otherwise + * successful result into the retry logic. + * + *

A result this plugin produced earlier is passed straight through: reflecting on a reflection + * would count a single tool failure twice. + */ + @Override + public Maybe> afterToolCallback( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + if (ToolFailureResponse.isReflection(result)) { + return Maybe.empty(); + } + return extractErrorFromResult(tool, toolArgs, toolContext, result) + .flatMap(error -> handleToolError(tool, toolArgs, toolContext, error)) + .switchIfEmpty(Maybe.fromRunnable(() -> resetFailures(tool, toolContext))); + } + + /** Turns a thrown tool error into reflection guidance for the model. */ + @Override + public Maybe> onToolErrorCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { + return handleToolError(tool, toolArgs, toolContext, error); + } + + /** + * Detects an error inside a tool result that did not throw — for example {@code + * {"status": "error"}} — so it can drive the same retry logic. + * + *

Empty by default, exactly as in adk-python. Override to opt in. + */ + protected Maybe extractErrorFromResult( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + return Maybe.empty(); + } + + /** + * The key failure counts are grouped under. Override to track per user or per session instead of + * the configured {@link TrackingScope}. + */ + protected String scopeKey(ToolContext toolContext) { + return switch (trackingScope) { + case INVOCATION -> toolContext.invocationId(); + case GLOBAL -> GLOBAL_SCOPE_KEY; + }; + } + + /** + * Counts the failure and decides between guidance, a final message, or propagating the error. + * + *

Never completes empty. {@link #afterToolCallback} treats an empty result as "the tool + * succeeded" and resets the counter, so an empty return here would clear the count of the very + * call that just failed. + */ + private Maybe> handleToolError( + BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { + if (maxRetries == 0) { + return exhausted(tool, toolArgs, error, FIRST_FAILURE); + } + int attempt = failures.recordFailure(scopeKey(toolContext), tool.name()); + if (attempt <= maxRetries) { + return Maybe.just(reflection(tool, toolArgs, error, attempt)); + } + return exhausted(tool, toolArgs, error, attempt); + } + + /** + * Either propagates the final error or hands back the give-up message, per configuration. + * + *

{@code failures} is the number of consecutive failures actually observed, which is what both + * the give-up message and the response's {@code retry_count} report. Retrying is disabled at + * {@code maxRetries == 0}, where nothing is counted at all — upstream returns before its counter + * runs ({@code reflect_retry_tool_plugin.py:243-246}) and so does this — so the observed count + * there is the one failure being reported. + */ + private Maybe> exhausted( + BaseTool tool, Map toolArgs, Throwable error, int failures) { + return throwExceptionIfRetryExceeded + ? Maybe.error(error) + : Maybe.just(retryExceeded(tool, toolArgs, error, failures)); + } + + private void resetFailures(BaseTool tool, ToolContext toolContext) { + failures.reset(scopeKey(toolContext), tool.name()); + } + + private ImmutableMap reflection( + BaseTool tool, Map toolArgs, Throwable error, int attempt) { + return response( + error, attempt, ReflectionGuidance.forRetry(tool, toolArgs, error, attempt, maxRetries)); + } + + private ImmutableMap retryExceeded( + BaseTool tool, Map toolArgs, Throwable error, int failures) { + return response( + error, failures, ReflectionGuidance.forExhausted(tool, toolArgs, error, failures)); + } + + private static ImmutableMap response( + Throwable error, int retryCount, String guidance) { + return new ToolFailureResponse( + error.getClass().getSimpleName(), + Strings.nullToEmpty(error.getMessage()), + retryCount, + guidance) + .toMap(); + } +} diff --git a/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectionGuidance.java b/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectionGuidance.java new file mode 100644 index 000000000..08c2c6a6a --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/reflectandretry/ReflectionGuidance.java @@ -0,0 +1,130 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.adk.JsonBaseModel; +import com.google.adk.tools.BaseTool; +import com.google.common.base.Strings; +import java.util.Map; + +/** + * The text {@link ReflectAndRetryToolPlugin} sends the model after a tool fails. + * + *

Kept apart from the plugin so that one class decides whether to retry and this one + * decides only what the model is told. Both messages are ports of the f-strings in + * adk-python's {@code reflect_retry_tool_plugin.py} and are reproduced verbatim. + */ +final class ReflectionGuidance { + + private static final String ERROR_DETAILS_FORMAT = "%s: %s"; + + private static final String RETRY = + """ + The call to tool `%s` failed. + + **Error Details:** + ``` + %s + ``` + + **Tool Arguments Used:** + ```json + %s + ``` + + **Reflection Guidance:** + This is retry attempt **%d of %d**. Analyze the error and the arguments you provided. Do not \ + repeat the exact same call. Consider the following before your next attempt: + + 1. **Invalid Parameters**: Does the error suggest that one or more arguments are incorrect, \ + badly formatted, or missing? Review the tool's schema and your arguments. + 2. **State or Preconditions**: Did a previous step fail or not produce the necessary \ + state/resource for this tool to succeed? + 3. **Alternative Approach**: Is this the right tool for the job? Could another tool or a \ + different sequence of steps achieve the goal? + 4. **Simplify the Task**: Can you break the problem down into smaller, simpler steps? + 5. **Wrong Function Name**: Does the error indicates the tool is not found? Please check \ + again and only use available tools. + + Formulate a new plan based on your analysis and try a corrected or different approach."""; + + private static final String EXHAUSTED = + """ + The tool `%s` has failed consecutively %d times and the retry limit has been exceeded. + + **Last Error:** + ``` + %s + ``` + + **Last Arguments Used:** + ```json + %s + ``` + + **Final Instruction:** + **Do not attempt to use the `%s` tool again for this task.** You must now try a different \ + approach. Acknowledge the failure and devise a new strategy, potentially using other \ + available tools or informing the user that the task cannot be completed."""; + + private ReflectionGuidance() {} + + /** Asks the model to analyze the failure and try a corrected call. */ + static String forRetry( + BaseTool tool, Map toolArgs, Throwable error, int attempt, int maxRetries) { + return RETRY.formatted( + tool.name(), errorDetails(error), argsAsJson(toolArgs), attempt, maxRetries); + } + + /** + * Tells the model to stop calling the tool and change approach. + * + *

{@code failures} is the number of consecutive failures observed. adk-python interpolates its + * configured {@code max_retries} here instead ({@code reflect_retry_tool_plugin.py:357}), which + * under-reports by one at every setting because the give-up fires on the failure *after* the + * limit — and at {@code max_retries=0} tells the model the tool "has failed consecutively 0 times + * and the retry limit has been exceeded". This port reports what actually happened. + */ + static String forExhausted( + BaseTool tool, Map toolArgs, Throwable error, int failures) { + return EXHAUSTED.formatted( + tool.name(), failures, errorDetails(error), argsAsJson(toolArgs), tool.name()); + } + + private static String errorDetails(Throwable error) { + return ERROR_DETAILS_FORMAT.formatted( + error.getClass().getSimpleName(), Strings.nullToEmpty(error.getMessage())); + } + + /** + * Pretty-prints the arguments for the guidance message, falling back to the map's own rendering. + * + *

Never throws: a serialization failure must not mask the tool failure being reported, and the + * model still needs the echo of the arguments it sent. Mirrors adk-python's {@code + * json.dumps(..., default=str)}. + */ + private static String argsAsJson(Map toolArgs) { + try { + return JsonBaseModel.getMapper() + .writerWithDefaultPrettyPrinter() + .writeValueAsString(toolArgs); + } catch (JsonProcessingException e) { + return String.valueOf(toolArgs); + } + } +} diff --git a/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureResponse.java b/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureResponse.java new file mode 100644 index 000000000..3dedd4203 --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureResponse.java @@ -0,0 +1,58 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import com.google.common.collect.ImmutableMap; +import java.util.Map; + +/** + * The response {@link ReflectAndRetryToolPlugin} substitutes for a failed tool call: what went + * wrong, how many times it has now failed, and what the model should do differently. + * + *

Port of adk-python's {@code ToolFailureResponse} pydantic model. The map keys are the + * snake_case field names that model serializes to; {@link #RESPONSE_TYPE_KEY} is how the plugin + * recognizes its own output on a later {@code afterToolCallback}. + */ +record ToolFailureResponse( + String errorType, String errorDetails, int retryCount, String reflectionGuidance) { + + /** Marks a tool result as this plugin's own output rather than a real tool response. */ + static final String REFLECT_AND_RETRY_RESPONSE_TYPE = "ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN"; + + static final String RESPONSE_TYPE_KEY = "response_type"; + + private static final String ERROR_TYPE_KEY = "error_type"; + private static final String ERROR_DETAILS_KEY = "error_details"; + private static final String RETRY_COUNT_KEY = "retry_count"; + private static final String REFLECTION_GUIDANCE_KEY = "reflection_guidance"; + + /** The tool-response map handed back to the flow, matching adk-python's serialization. */ + ImmutableMap toMap() { + return ImmutableMap.of( + RESPONSE_TYPE_KEY, REFLECT_AND_RETRY_RESPONSE_TYPE, + ERROR_TYPE_KEY, errorType, + ERROR_DETAILS_KEY, errorDetails, + RETRY_COUNT_KEY, retryCount, + REFLECTION_GUIDANCE_KEY, reflectionGuidance); + } + + /** Whether {@code result} is a response this plugin produced earlier. */ + static boolean isReflection(Object result) { + return result instanceof Map map + && REFLECT_AND_RETRY_RESPONSE_TYPE.equals(map.get(RESPONSE_TYPE_KEY)); + } +} diff --git a/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureTracker.java b/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureTracker.java new file mode 100644 index 000000000..ca5b81396 --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/reflectandretry/ToolFailureTracker.java @@ -0,0 +1,53 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Consecutive-failure counts, one per tool per tracking scope. + * + *

Kept apart from {@link ReflectAndRetryToolPlugin} so the plugin decides what a failure means + * and this class decides only how the counts are held. + * + *

adk-python nests a {@code dict} per scope inside another {@code dict} and guards both with an + * {@code asyncio.Lock}, because it mutates them across {@code await} points. Neither is needed + * here. Nothing ever enumerates the tools within one scope, so the pair is simply the key, and + * {@link Map#merge} and {@link Map#remove} each apply atomically on a {@link ConcurrentHashMap} — + * parallel tool failures cannot lose a count. + */ +final class ToolFailureTracker { + + private final ConcurrentHashMap counts = new ConcurrentHashMap<>(); + + /** Records one more consecutive failure and returns the new count, atomically. */ + int recordFailure(String scopeKey, String toolName) { + return counts.merge(new ToolInScope(scopeKey, toolName), 1, Integer::sum); + } + + /** + * Clears one tool's failure count, leaving every other tool's untouched — a success with one tool + * must not forgive another's failures. + */ + void reset(String scopeKey, String toolName) { + counts.remove(new ToolInScope(scopeKey, toolName)); + } + + /** The counter key: one tool, within one tracking scope. */ + private record ToolInScope(String scopeKey, String toolName) {} +} diff --git a/core/src/main/java/com/google/adk/plugins/reflectandretry/TrackingScope.java b/core/src/main/java/com/google/adk/plugins/reflectandretry/TrackingScope.java new file mode 100644 index 000000000..da09c1e0f --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/reflectandretry/TrackingScope.java @@ -0,0 +1,30 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +/** + * Defines the lifecycle scope for tracking tool failure counts. + * + *

Port of {@code TrackingScope} in adk-python's {@code reflect_retry_tool_plugin.py}. + */ +public enum TrackingScope { + /** Failure counts live for one invocation and are keyed by its invocation id. */ + INVOCATION, + + /** Failure counts live for the lifetime of the plugin instance, shared across invocations. */ + GLOBAL; +} diff --git a/core/src/test/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPluginTest.java b/core/src/test/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPluginTest.java new file mode 100644 index 000000000..16d10c4f3 --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/reflectandretry/ReflectAndRetryToolPluginTest.java @@ -0,0 +1,331 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.when; + +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableMap; +import io.reactivex.rxjava3.core.Maybe; +import java.util.Map; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; + +@RunWith(JUnit4.class) +public class ReflectAndRetryToolPluginTest { + + private static final String INVOCATION_ID = "invocation-1"; + private static final String OTHER_INVOCATION_ID = "invocation-2"; + private static final String USER_ID = "user-42"; + private static final String TOOL_NAME = "flaky_tool"; + private static final String OTHER_TOOL_NAME = "other_tool"; + private static final ImmutableMap ARGS = ImmutableMap.of("city", "Phoenix"); + + @Rule public MockitoRule mockitoRule = MockitoJUnit.rule(); + + @Mock private BaseTool mockTool; + @Mock private BaseTool mockOtherTool; + @Mock private ToolContext mockToolContext; + + private final RuntimeException error = new IllegalStateException("boom"); + + @Before + public void setUp() { + when(mockTool.name()).thenReturn(TOOL_NAME); + when(mockToolContext.invocationId()).thenReturn(INVOCATION_ID); + } + + @Test + public void onToolError_withinRetryLimit_returnsReflectionGuidance() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + + Map response = callOnError(plugin); + + assertThat(response) + .containsEntry("response_type", "ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN"); + assertThat(response).containsEntry("error_type", "IllegalStateException"); + assertThat(response).containsEntry("error_details", "boom"); + assertThat(response).containsEntry("retry_count", 1); + assertThat((String) response.get("reflection_guidance")).contains("retry attempt **1 of 3**"); + } + + @Test + public void onToolError_guidanceCarriesToolNameAndArguments() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + + String guidance = (String) callOnError(plugin).get("reflection_guidance"); + + assertThat(guidance).contains(TOOL_NAME); + assertThat(guidance).contains("\"city\" : \"Phoenix\""); + assertThat(guidance).contains("IllegalStateException: boom"); + } + + @Test + public void onToolError_errorWithoutMessage_reportsEmptyDetails() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + + Map response = callOnError(plugin, new IllegalStateException()); + + assertThat(response).containsEntry("error_details", ""); + assertThat((String) response.get("reflection_guidance")) + .contains("```\nIllegalStateException: \n```"); + } + + @Test + public void onToolError_pastRetryLimit_errorWithoutMessage_reportsEmptyDetails() { + ReflectAndRetryToolPlugin plugin = + new ReflectAndRetryToolPlugin("p", 1, false, TrackingScope.INVOCATION); + IllegalStateException messageless = new IllegalStateException(); + + callOnError(plugin, messageless); + Map exceeded = callOnError(plugin, messageless); + + assertThat(exceeded).containsEntry("error_details", ""); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("```\nIllegalStateException: \n```"); + } + + @Test + public void onToolError_countsConsecutiveFailures() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + + callOnError(plugin); + callOnError(plugin); + Map third = callOnError(plugin); + + assertThat(third).containsEntry("retry_count", 3); + } + + @Test + public void onToolError_pastRetryLimit_propagatesOriginalError() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(2); + + callOnError(plugin); + callOnError(plugin); + + plugin.onToolErrorCallback(mockTool, ARGS, mockToolContext, error).test().assertError(error); + } + + @Test + public void onToolError_pastRetryLimit_whenNotThrowing_returnsGiveUpMessage() { + ReflectAndRetryToolPlugin plugin = + new ReflectAndRetryToolPlugin("p", 1, false, TrackingScope.INVOCATION); + + callOnError(plugin); + Map exceeded = callOnError(plugin); + + assertThat((String) exceeded.get("reflection_guidance")) + .contains("Do not attempt to use the `flaky_tool` tool again"); + assertThat(exceeded).containsEntry("retry_count", 2); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("has failed consecutively 2 times"); + } + + @Test + public void maxRetriesZero_propagatesImmediately() { + ReflectAndRetryToolPlugin plugin = + new ReflectAndRetryToolPlugin("p", 0, true, TrackingScope.INVOCATION); + + plugin.onToolErrorCallback(mockTool, ARGS, mockToolContext, error).test().assertError(error); + } + + @Test + public void maxRetriesZero_whenNotThrowing_returnsGiveUpMessage() { + ReflectAndRetryToolPlugin plugin = + new ReflectAndRetryToolPlugin("p", 0, false, TrackingScope.INVOCATION); + + Map exceeded = callOnError(plugin); + + assertThat(exceeded).containsEntry("retry_count", 1); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("has failed consecutively 1 times"); + } + + @Test + public void onToolError_whenArgumentsCannotBeSerialized_fallsBackToMapRendering() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + ImmutableMap unserializable = ImmutableMap.of("payload", new Opaque()); + + Map response = + plugin.onToolErrorCallback(mockTool, unserializable, mockToolContext, error).blockingGet(); + + assertThat((String) response.get("reflection_guidance")).contains("{payload=}"); + } + + @Test + public void negativeMaxRetries_isRejected() { + assertThrows( + IllegalArgumentException.class, + () -> new ReflectAndRetryToolPlugin("p", -1, true, TrackingScope.INVOCATION)); + } + + @Test + public void nullTrackingScope_isRejectedAtConstruction() { + assertThrows( + NullPointerException.class, () -> new ReflectAndRetryToolPlugin("p", 3, true, null)); + } + + @Test + public void afterTool_onSuccess_resetsThatToolsCounter() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + callOnError(plugin); + callOnError(plugin); + + plugin + .afterToolCallback(mockTool, ARGS, mockToolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterTool_successOfOneTool_doesNotForgiveAnother() { + when(mockOtherTool.name()).thenReturn(OTHER_TOOL_NAME); + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + callOnError(plugin); + + plugin + .afterToolCallback(mockOtherTool, ARGS, mockToolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + @Test + public void afterTool_ownReflectionResponse_isPassedThroughUncounted() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + Map reflection = callOnError(plugin); + + plugin + .afterToolCallback(mockTool, ARGS, mockToolContext, reflection) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + @Test + public void invocationScope_countsAreIsolatedPerInvocation() { + ReflectAndRetryToolPlugin plugin = new ReflectAndRetryToolPlugin(3); + callOnError(plugin); + + when(mockToolContext.invocationId()).thenReturn(OTHER_INVOCATION_ID); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void globalScope_countsSurviveAcrossInvocations() { + ReflectAndRetryToolPlugin plugin = + new ReflectAndRetryToolPlugin("p", 3, true, TrackingScope.GLOBAL); + callOnError(plugin); + + when(mockToolContext.invocationId()).thenReturn(OTHER_INVOCATION_ID); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + @Test + public void extractErrorFromResult_whenOverridden_treatsSuccessfulResultAsFailure() { + ReflectAndRetryToolPlugin plugin = new StatusAwarePlugin(); + + Map response = + plugin + .afterToolCallback(mockTool, ARGS, mockToolContext, ImmutableMap.of("status", "error")) + .blockingGet(); + + assertThat(response).containsEntry("retry_count", 1); + assertThat(response).containsEntry("error_type", "IllegalStateException"); + } + + @Test + public void scopeKey_whenOverridden_outranksTheConfiguredTrackingScope() { + ReflectAndRetryToolPlugin plugin = new PerUserPlugin(); + callOnError(plugin); + + when(mockToolContext.invocationId()).thenReturn(OTHER_INVOCATION_ID); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + private Map callOnError(ReflectAndRetryToolPlugin plugin) { + return callOnError(plugin, error); + } + + private Map callOnError(ReflectAndRetryToolPlugin plugin, Throwable thrown) { + return plugin.onToolErrorCallback(mockTool, ARGS, mockToolContext, thrown).blockingGet(); + } + + /** No properties, so Jackson refuses to serialize it; only its {@code toString} can be used. */ + private static final class Opaque { + + @Override + public String toString() { + return ""; + } + } + + /** + * Tracks failures per user, the documented reason to override {@code scopeKey}. Configured with + * {@link TrackingScope#INVOCATION} so that counts surviving a change of invocation id can only be + * the override taking effect. + */ + private static final class PerUserPlugin extends ReflectAndRetryToolPlugin { + + PerUserPlugin() { + super("per_user", 3, true, TrackingScope.INVOCATION); + } + + @Override + protected String scopeKey(ToolContext toolContext) { + return USER_ID; + } + } + + /** Treats {@code {"status": "error"}} as a failure, the documented reason to override. */ + private static final class StatusAwarePlugin extends ReflectAndRetryToolPlugin { + + StatusAwarePlugin() { + super("status_aware", 3, true, TrackingScope.INVOCATION); + } + + @Override + protected Maybe extractErrorFromResult( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + return "error".equals(result.get("status")) + ? Maybe.just(new IllegalStateException("tool reported status=error")) + : Maybe.empty(); + } + } +} diff --git a/core/src/test/java/com/google/adk/plugins/reflectandretry/ToolFailureTrackerTest.java b/core/src/test/java/com/google/adk/plugins/reflectandretry/ToolFailureTrackerTest.java new file mode 100644 index 000000000..1cb85c55e --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/reflectandretry/ToolFailureTrackerTest.java @@ -0,0 +1,150 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.reflectandretry; + +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.Callable; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.stream.IntStream; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Pins the guarantee {@link ToolFailureTracker} makes in place of adk-python's {@code + * asyncio.Lock}: parallel failures of one tool are each counted exactly once. + */ +@RunWith(JUnit4.class) +public class ToolFailureTrackerTest { + + private static final String SCOPE = "invocation-1"; + private static final String OTHER_SCOPE = "invocation-2"; + private static final String TOOL_NAME = "flaky_tool"; + private static final String OTHER_TOOL_NAME = "other_tool"; + private static final int THREADS = 8; + private static final int FAILURES_PER_THREAD = 2_000; + private static final int TOTAL_FAILURES = THREADS * FAILURES_PER_THREAD; + private static final int TIMEOUT_SECONDS = 30; + + /** + * Every thread hammers the same tool's counter from a common start. A lost update repeats a count + * and a dropped one skips it, so the counts handed back across all threads are exactly {@code + * 1..n} if and only if every update was atomic. + * + *

Contention has to be sustained to be meaningful: with one call per thread the + * critical section is far shorter than the spread in thread wake-up, and a plain {@link + * java.util.HashMap} passes. + */ + @Test + public void recordFailure_inParallel_countsEachFailureExactlyOnce() throws Exception { + ToolFailureTracker tracker = new ToolFailureTracker(); + CyclicBarrier startTogether = new CyclicBarrier(THREADS); + ExecutorService threads = Executors.newFixedThreadPool(THREADS); + + List>> counts = + threads.invokeAll(Collections.nCopies(THREADS, recordRepeatedlyAt(tracker, startTogether))); + threads.shutdown(); + + assertThat(threads.awaitTermination(TIMEOUT_SECONDS, TimeUnit.SECONDS)).isTrue(); + assertThat(resolve(counts)).containsExactlyElementsIn(oneTo(TOTAL_FAILURES)); + } + + @Test + public void recordFailure_countsEachToolSeparately() { + ToolFailureTracker tracker = new ToolFailureTracker(); + + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(SCOPE, OTHER_TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(2); + } + + @Test + public void recordFailure_countsEachScopeSeparately() { + ToolFailureTracker tracker = new ToolFailureTracker(); + + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(OTHER_SCOPE, TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(2); + } + + @Test + public void reset_clearsOnlyThatScope() { + ToolFailureTracker tracker = new ToolFailureTracker(); + tracker.recordFailure(SCOPE, TOOL_NAME); + tracker.recordFailure(OTHER_SCOPE, TOOL_NAME); + + tracker.reset(SCOPE, TOOL_NAME); + + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(OTHER_SCOPE, TOOL_NAME)).isEqualTo(2); + } + + @Test + public void reset_clearsOnlyThatTool() { + ToolFailureTracker tracker = new ToolFailureTracker(); + tracker.recordFailure(SCOPE, TOOL_NAME); + tracker.recordFailure(SCOPE, OTHER_TOOL_NAME); + + tracker.reset(SCOPE, TOOL_NAME); + + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(1); + assertThat(tracker.recordFailure(SCOPE, OTHER_TOOL_NAME)).isEqualTo(2); + } + + @Test + public void reset_onUnknownScope_isASilentNoOp() { + ToolFailureTracker tracker = new ToolFailureTracker(); + + tracker.reset("never-seen", TOOL_NAME); + + assertThat(tracker.recordFailure(SCOPE, TOOL_NAME)).isEqualTo(1); + } + + private static Callable> recordRepeatedlyAt( + ToolFailureTracker tracker, CyclicBarrier barrier) { + return () -> recordRepeatedlyAfterBarrier(tracker, barrier); + } + + private static ImmutableList recordRepeatedlyAfterBarrier( + ToolFailureTracker tracker, CyclicBarrier barrier) throws Exception { + barrier.await(TIMEOUT_SECONDS, TimeUnit.SECONDS); + return IntStream.range(0, FAILURES_PER_THREAD) + .mapToObj(unused -> tracker.recordFailure(SCOPE, TOOL_NAME)) + .collect(ImmutableList.toImmutableList()); + } + + private static ImmutableList resolve(List>> futures) + throws Exception { + ImmutableList.Builder counts = ImmutableList.builder(); + for (Future> future : futures) { + counts.addAll(future.get()); + } + return counts.build(); + } + + private static ImmutableList oneTo(int last) { + return IntStream.rangeClosed(1, last).boxed().collect(ImmutableList.toImmutableList()); + } +}