From 33e9f8b2dcd88916d12621243fa6e0b1988c5bad Mon Sep 17 00:00:00 2001 From: svetanis Date: Sun, 2 Aug 2026 14:05:51 -0700 Subject: [PATCH] feat(plugins): port ReflectAndRetryToolPlugin from adk-python --- .../plugins/ReflectAndRetryToolPlugin.java | 519 ++++++++++ .../ReflectAndRetryToolPluginTest.java | 963 ++++++++++++++++++ 2 files changed, 1482 insertions(+) create mode 100644 core/src/main/java/com/google/adk/plugins/ReflectAndRetryToolPlugin.java create mode 100644 core/src/test/java/com/google/adk/plugins/ReflectAndRetryToolPluginTest.java diff --git a/core/src/main/java/com/google/adk/plugins/ReflectAndRetryToolPlugin.java b/core/src/main/java/com/google/adk/plugins/ReflectAndRetryToolPlugin.java new file mode 100644 index 000000000..3c675b230 --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/ReflectAndRetryToolPlugin.java @@ -0,0 +1,519 @@ +/* + * 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; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.adk.JsonBaseModel; +import com.google.adk.agents.InvocationContext; +import com.google.adk.annotations.Experimental; +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 com.google.errorprone.annotations.CanIgnoreReturnValue; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * 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. + * + *

A failure is an error a tool emits, such as {@code Single.error}, or an error the {@link + * ResultErrorExtractor} finds in a result. {@code FunctionTool} returns an exception thrown by its + * method as a {@code {"status": "error"}} result, so such failures reach this plugin only through a + * result error extractor. Calls to a tool name the agent does not have never reach this plugin. + * + *

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

Example: + * + * 

{@code
+ * Runner runner =
+ *     new InMemoryRunner(
+ *         agent,
+ *         APP_NAME,
+ *         ImmutableList.of(ReflectAndRetryToolPlugin.builder().maxRetries(3).build()));
+ * }
+ */ +@Experimental +public final 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 {@code error_type} reported for an error that is not a {@link Throwable}. */ + private static final String TOOL_ERROR_TYPE = "ToolError"; + + private static final String RESULT_ERROR_EXCEEDED = + "A tool result reported an error and the retry limit was exceeded."; + + private static final ResultErrorExtractor NO_RESULT_ERRORS = + (tool, toolArgs, toolContext, result) -> Optional.empty(); + + /** Marks a tool result as this plugin's own output rather than a real tool response. */ + private static final String REFLECT_AND_RETRY_RESPONSE_TYPE = + "ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN"; + + private 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"; + + private static final String ERROR_DETAILS_FORMAT = "%s: %s"; + + /** Written in place of a map or list that contains itself. */ + private static final String CIRCULAR_REFERENCE = ""; + + private static final String RETRY_GUIDANCE = + """ + 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_GUIDANCE = + """ + 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 final int maxRetries; + private final boolean throwExceptionIfRetryExceeded; + private final TrackingScope trackingScope; + private final ResultErrorExtractor resultErrorExtractor; + + /** Consecutive-failure counts, one per tool per tracking scope, updated atomically. */ + private final ConcurrentHashMap failureCounts = new ConcurrentHashMap<>(); + + /** Defines the lifecycle scope for tracking tool failure counts. */ + @Experimental + 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; + } + + /** Finds an error that a tool reported in its result instead of throwing. */ + @Experimental + @FunctionalInterface + public interface ResultErrorExtractor { + + /** + * Returns the error in {@code result}, such as {@code result} itself when it holds {@code + * "status": "error"}, or empty if the call succeeded. + * + *

An error that is not a {@link Throwable} is reported to the model with {@code error_type} + * {@code "ToolError"}; a {@link Map} error is passed as data, so like a tool result it must + * hold only JSON-native values. Past the limit, when {@code throwExceptionIfRetryExceeded} is + * set, a {@link Throwable} error is emitted as is and any other error as an {@link + * IllegalStateException} with a generic message. + */ + Optional extractErrorFromResult( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result); + } + + private ReflectAndRetryToolPlugin(Builder builder) { + super(builder.name); + this.maxRetries = builder.maxRetries; + this.throwExceptionIfRetryExceeded = builder.throwExceptionIfRetryExceeded; + this.trackingScope = builder.trackingScope; + this.resultErrorExtractor = builder.resultErrorExtractor; + } + + /** + * Returns a builder defaulting to three retries, throwing once they are exceeded, counted per + * invocation, with no result errors. + */ + public static Builder builder() { + return new Builder(); + } + + /** Builder for {@link ReflectAndRetryToolPlugin}. */ + public static final class Builder { + private int maxRetries = DEFAULT_MAX_RETRIES; + private boolean throwExceptionIfRetryExceeded = true; + private TrackingScope trackingScope = TrackingScope.INVOCATION; + private ResultErrorExtractor resultErrorExtractor = NO_RESULT_ERRORS; + private String name = DEFAULT_NAME; + + private Builder() {} + + /** + * Sets the number of consecutive failures of one tool answered with retry guidance; {@code 0} + * disables retrying. + * + * @throws IllegalArgumentException if {@code maxRetries} is negative + */ + @CanIgnoreReturnValue + public Builder maxRetries(int maxRetries) { + checkArgument(maxRetries >= 0, NEGATIVE_RETRIES, maxRetries); + this.maxRetries = maxRetries; + return this; + } + + /** + * Sets whether, past the limit, the error propagates rather than the model being told to stop + * using the tool. + */ + @CanIgnoreReturnValue + public Builder throwExceptionIfRetryExceeded(boolean throwExceptionIfRetryExceeded) { + this.throwExceptionIfRetryExceeded = throwExceptionIfRetryExceeded; + return this; + } + + /** Sets whether failures are counted per invocation or across invocations. */ + @CanIgnoreReturnValue + public Builder trackingScope(TrackingScope trackingScope) { + this.trackingScope = checkNotNull(trackingScope); + return this; + } + + /** Sets how errors reported inside tool results are found. */ + @CanIgnoreReturnValue + public Builder resultErrorExtractor(ResultErrorExtractor resultErrorExtractor) { + this.resultErrorExtractor = checkNotNull(resultErrorExtractor); + return this; + } + + /** Sets the plugin instance name. */ + @CanIgnoreReturnValue + public Builder name(String name) { + this.name = checkNotNull(name); + return this; + } + + public ReflectAndRetryToolPlugin build() { + return new ReflectAndRetryToolPlugin(this); + } + } + + /** + * Resets the tool's failure count on success, or routes an error carried inside an otherwise + * successful result into the retry logic. A response this plugin produced is passed through. + */ + @Override + public Maybe> afterToolCallback( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + if (isReflection(result)) { + return Maybe.empty(); + } + return Maybe.defer(() -> checkResult(tool, toolArgs, toolContext, result)); + } + + /** Drops the ended invocation's failure counts. */ + @Override + public Completable afterRunCallback(InvocationContext invocationContext) { + return Completable.fromAction(() -> clearInvocation(invocationContext)); + } + + /** Drops the failed invocation's failure counts. */ + @Override + public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) { + return Completable.fromAction(() -> clearInvocation(invocationContext)); + } + + /** Turns a thrown tool error into reflection guidance for the model. */ + @Override + public Maybe> onToolErrorCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { + return Maybe.defer(() -> handleToolError(tool, toolArgs, toolContext, error)); + } + + /** The key failure counts are grouped under, per the configured {@link TrackingScope}. */ + private String scopeKey(ToolContext toolContext) { + return switch (trackingScope) { + case INVOCATION -> toolContext.invocationId(); + case GLOBAL -> GLOBAL_SCOPE_KEY; + }; + } + + /** Resets the tool's count if its result reports no error, else handles the reported error. */ + private Maybe> checkResult( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + Optional error = + resultErrorExtractor.extractErrorFromResult(tool, toolArgs, toolContext, result); + if (error.isEmpty()) { + resetFailures(tool, toolContext); + return Maybe.empty(); + } + return handleResultError(tool, toolArgs, toolContext, error.get()); + } + + /** + * Handles an error found inside a tool result, emitting the error where {@link #handleToolError} + * completes empty. + */ + private Maybe> handleResultError( + BaseTool tool, Map toolArgs, ToolContext toolContext, Object error) { + return handleToolError(tool, toolArgs, toolContext, error) + .switchIfEmpty(Maybe.error(asThrowable(error))); + } + + /** + * Counts the failure and returns guidance or the give-up message. Completes empty when the limit + * is reached and the original error should propagate. + */ + private Maybe> handleToolError( + BaseTool tool, Map toolArgs, ToolContext toolContext, Object error) { + if (maxRetries == 0) { + return exhausted(tool, toolArgs, error); + } + int attempt = recordFailure(scopeKey(toolContext), tool.name()); + if (attempt <= maxRetries) { + return Maybe.just(reflection(tool, toolArgs, error, attempt)); + } + return exhausted(tool, toolArgs, error); + } + + /** Completes empty to let the error propagate, or hands back the give-up message. */ + private Maybe> exhausted( + BaseTool tool, Map toolArgs, Object error) { + return throwExceptionIfRetryExceeded + ? Maybe.empty() + : Maybe.just(retryExceeded(tool, toolArgs, error)); + } + + /** Records one more consecutive failure and returns the new count, atomically. */ + private int recordFailure(String scopeKey, String toolName) { + return failureCounts.merge(new ToolInScope(scopeKey, toolName), 1, Integer::sum); + } + + /** Clears one tool's failure count within its scope, leaving every other tool's untouched. */ + private void resetFailures(BaseTool tool, ToolContext toolContext) { + failureCounts.remove(new ToolInScope(scopeKey(toolContext), tool.name())); + } + + /** + * Removes every count keyed by the invocation's id. Counts kept under any other key, such as + * {@link TrackingScope#GLOBAL}'s, are untouched. + */ + private void clearInvocation(InvocationContext invocationContext) { + String invocationId = invocationContext.invocationId(); + failureCounts.keySet().removeIf(key -> key.scopeKey().equals(invocationId)); + } + + private ImmutableMap reflection( + BaseTool tool, Map toolArgs, Object error, int attempt) { + return response(error, attempt, retryGuidance(tool, toolArgs, error, attempt)); + } + + private ImmutableMap retryExceeded( + BaseTool tool, Map toolArgs, Object error) { + return response(error, maxRetries, exhaustedGuidance(tool, toolArgs, error)); + } + + /** Asks the model to analyze the failure and try a corrected call. */ + private String retryGuidance( + BaseTool tool, Map toolArgs, Object error, int attempt) { + return RETRY_GUIDANCE.formatted( + tool.name(), formatErrorDetails(error), asJson(toolArgs), attempt, maxRetries); + } + + /** Tells the model to stop calling the tool and change approach. */ + private String exhaustedGuidance(BaseTool tool, Map toolArgs, Object error) { + return EXHAUSTED_GUIDANCE.formatted( + tool.name(), maxRetries, formatErrorDetails(error), asJson(toolArgs), tool.name()); + } + + /** The error as shown inside the guidance text. */ + private static String formatErrorDetails(Object error) { + if (error instanceof Throwable throwable) { + return ERROR_DETAILS_FORMAT.formatted( + errorType(throwable), Strings.nullToEmpty(throwable.getMessage())); + } + return error instanceof Map map ? asJson(map) : String.valueOf(error); + } + + /** + * Pretty-prints a map as JSON for the guidance message. A value that cannot be serialized is + * written as its string form. + */ + private static String asJson(Map map) { + try { + return writeJson(map); + } catch (JsonProcessingException e) { + return asJsonWithTextFallback(map); + } + } + + private static String asJsonWithTextFallback(Map map) { + try { + return writeJson(jsonSafe(map, Collections.newSetFromMap(new IdentityHashMap<>()))); + } catch (JsonProcessingException e) { + return String.valueOf(map); + } + } + + /** + * Copies maps, lists and arrays, replacing each leaf that cannot be serialized with its string + * form, and a container that contains itself with {@link #CIRCULAR_REFERENCE}. + * + * @param path the containers currently being copied, compared by identity + */ + private static Object jsonSafe(Object value, Set path) { + if (!isContainer(value)) { + return isSerializable(value) ? value : String.valueOf(value); + } + if (!path.add(value)) { + return CIRCULAR_REFERENCE; + } + Object copy = + value instanceof Map map ? jsonSafeMap(map, path) : jsonSafeList(items(value), path); + path.remove(value); + return copy; + } + + private static boolean isContainer(Object value) { + return value instanceof Map || value instanceof Iterable || value instanceof Object[]; + } + + private static Iterable items(Object container) { + return container instanceof Object[] array ? Arrays.asList(array) : (Iterable) container; + } + + private static Map jsonSafeMap(Map map, Set path) { + Map copy = new LinkedHashMap<>(); + map.forEach((key, item) -> copy.put(String.valueOf(key), jsonSafe(item, path))); + return Collections.unmodifiableMap(copy); + } + + private static List jsonSafeList(Iterable items, Set path) { + List copy = new ArrayList<>(); + items.forEach(item -> copy.add(jsonSafe(item, path))); + return Collections.unmodifiableList(copy); + } + + private static boolean isSerializable(Object value) { + try { + JsonBaseModel.getMapper().writeValueAsString(value); + return true; + } catch (JsonProcessingException e) { + return false; + } + } + + private static String writeJson(Object value) throws JsonProcessingException { + return JsonBaseModel.getMapper().writerWithDefaultPrettyPrinter().writeValueAsString(value); + } + + /** The tool response substituted for a failed call. */ + private static ImmutableMap response( + Object error, int retryCount, String guidance) { + return ImmutableMap.of( + RESPONSE_TYPE_KEY, REFLECT_AND_RETRY_RESPONSE_TYPE, + ERROR_TYPE_KEY, errorType(error), + ERROR_DETAILS_KEY, errorDetails(error), + RETRY_COUNT_KEY, retryCount, + REFLECTION_GUIDANCE_KEY, guidance); + } + + private static String errorType(Object error) { + return error instanceof Throwable ? error.getClass().getSimpleName() : TOOL_ERROR_TYPE; + } + + /** A {@link Throwable}'s message, a {@link Map} error as data, anything else as text. */ + private static Object errorDetails(Object error) { + if (error instanceof Throwable throwable) { + return Strings.nullToEmpty(throwable.getMessage()); + } + return error instanceof Map ? error : String.valueOf(error); + } + + /** The error to emit past the limit; one that is not a {@link Throwable} is tool output. */ + private static Throwable asThrowable(Object error) { + return error instanceof Throwable throwable + ? throwable + : new IllegalStateException(RESULT_ERROR_EXCEEDED); + } + + /** Whether {@code result} is a response this plugin produced earlier. */ + private static boolean isReflection(Object result) { + return result instanceof Map map + && REFLECT_AND_RETRY_RESPONSE_TYPE.equals(map.get(RESPONSE_TYPE_KEY)); + } + + /** The counter key: one tool, within one tracking scope. */ + private record ToolInScope(String scopeKey, String toolName) {} +} diff --git a/core/src/test/java/com/google/adk/plugins/ReflectAndRetryToolPluginTest.java b/core/src/test/java/com/google/adk/plugins/ReflectAndRetryToolPluginTest.java new file mode 100644 index 000000000..38747bfd2 --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/ReflectAndRetryToolPluginTest.java @@ -0,0 +1,963 @@ +/* + * 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; + +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createRootAgent; +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.adk.plugins.ReflectAndRetryToolPlugin.TrackingScope; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.Session; +import com.google.adk.testing.TestLlm; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.UUID; +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; + +@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 BOOM = "boom"; + private static final String APP_NAME = "test"; + private static final String USER_ID = "user"; + private static final String THROWING_METHOD = "alwaysThrows"; + private static final String UNKNOWN_TOOL = "no_such_tool"; + private static final String CALL_ID = "call-%s"; + private static final String RESPONSE_TYPE_KEY = "response_type"; + private static final String RESPONSE_TYPE = "ERROR_HANDLED_BY_REFLECT_AND_RETRY_PLUGIN"; + private static final int RUN_TIMEOUT_SECONDS = 10; + private static final ImmutableMap FAIL = ImmutableMap.of("fail", true); + private static final ImmutableMap SUCCEED = ImmutableMap.of("fail", false); + private static final String STATUS_ERROR_MESSAGE = "tool reported status=error"; + private static final String GENERIC_EXCEEDED_MESSAGE = + "A tool result reported an error and the retry limit was exceeded."; + private static final ImmutableMap STATUS_ERROR_RESULT = + ImmutableMap.of("status", "error"); + 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"); + 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 PARALLEL_TIMEOUT_SECONDS = 30; + + private final BaseTool tool = new FlakyTool(TOOL_NAME); + private final BaseTool otherTool = new FlakyTool(OTHER_TOOL_NAME); + private final InvocationContext invocation = invocation(INVOCATION_ID); + private final InvocationContext otherInvocation = invocation(OTHER_INVOCATION_ID); + private final ToolContext toolContext = ToolContext.builder(invocation).build(); + private final ToolContext otherToolContext = ToolContext.builder(otherInvocation).build(); + private final RuntimeException error = new IllegalStateException(BOOM); + + @Test + public void onToolError_withinRetryLimit_returnsReflectionGuidance() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + + Map response = callOnError(plugin); + + assertThat(response).containsEntry(RESPONSE_TYPE_KEY, RESPONSE_TYPE); + 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 = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + + 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 = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + + 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 = + ReflectAndRetryToolPlugin.builder() + .maxRetries(1) + .throwExceptionIfRetryExceeded(false) + .build(); + 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_countsNothingUntilSubscribed() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + + Maybe> failure = plugin.onToolErrorCallback(tool, ARGS, toolContext, error); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + + assertThat(failure.blockingGet()).containsEntry("retry_count", 2); + } + + @Test + public void afterTool_resetsNothingUntilSubscribed() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnError(plugin); + callOnError(plugin); + + Maybe> success = + plugin.afterToolCallback(tool, ARGS, toolContext, ImmutableMap.of("ok", true)); + assertThat(callOnError(plugin)).containsEntry("retry_count", 3); + + success.test().assertNoValues().assertComplete(); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterRun_clearsNothingUntilSubscribed() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnError(plugin); + + Completable cleanup = plugin.afterRunCallback(invocation); + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + + cleanup.test().assertComplete(); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void onRunError_clearsNothingUntilSubscribed() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnError(plugin); + + Completable cleanup = plugin.onRunErrorCallback(invocation, error); + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + + cleanup.test().assertComplete(); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void onToolError_countsConsecutiveFailures() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + + callOnError(plugin); + callOnError(plugin); + Map third = callOnError(plugin); + + assertThat(third).containsEntry("retry_count", 3); + } + + @Test + public void onToolError_pastRetryLimit_completesEmptySoTheOriginalErrorPropagates() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(2).build(); + + callOnError(plugin); + callOnError(plugin); + + plugin + .onToolErrorCallback(tool, ARGS, toolContext, error) + .test() + .assertNoValues() + .assertComplete(); + } + + @Test + public void onToolError_pastRetryLimit_whenNotThrowing_returnsGiveUpMessage() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(1) + .throwExceptionIfRetryExceeded(false) + .build(); + + 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", 1); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("has failed consecutively 1 times"); + } + + @Test + public void maxRetriesZero_completesEmptyImmediately() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(0).build(); + + plugin + .onToolErrorCallback(tool, ARGS, toolContext, error) + .test() + .assertNoValues() + .assertComplete(); + } + + @Test + public void maxRetriesZero_whenNotThrowing_returnsGiveUpMessage() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(0) + .throwExceptionIfRetryExceeded(false) + .build(); + + Map exceeded = callOnError(plugin); + + assertThat(exceeded).containsEntry("retry_count", 0); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("has failed consecutively 0 times"); + } + + @Test + public void onToolError_wellPastRetryLimit_keepsReportingMaxRetries() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(2) + .throwExceptionIfRetryExceeded(false) + .build(); + + callOnError(plugin); + callOnError(plugin); + callOnError(plugin); + Map fourth = callOnError(plugin); + + assertThat(fourth).containsEntry("retry_count", 2); + assertThat((String) fourth.get("reflection_guidance")) + .contains("has failed consecutively 2 times"); + } + + @Test + public void onToolError_whenAnArgumentCannotBeSerialized_writesOnlyThatValueAsText() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + ImmutableMap args = ImmutableMap.of("city", "Phoenix", "payload", new Opaque()); + + String guidance = guidance(plugin, args); + + assertThat(guidance).contains("\"city\" : \"Phoenix\""); + assertThat(guidance).contains("\"payload\" : \"\""); + } + + @Test + public void onToolError_whenAnArrayElementCannotBeSerialized_writesOnlyThatElementAsText() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + ImmutableMap args = ImmutableMap.of("items", new Object[] {1, new Opaque()}); + + String guidance = guidance(plugin, args); + + assertThat(guidance).containsMatch("\"items\" : \\[ 1, \"\" \\]"); + } + + @Test + public void onToolError_whenArgumentsContainThemselves_marksTheCycleInsteadOfThrowing() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + Map args = new HashMap<>(); + args.put("self", args); + + String guidance = guidance(plugin, args); + + assertThat(guidance).contains("\"self\" : \"\""); + } + + @Test + public void onToolError_whenAnArgumentAppearsTwice_writesItBothTimes() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + ImmutableMap shared = ImmutableMap.of("unit", "F"); + ImmutableMap args = + ImmutableMap.of("first", shared, "second", shared, "payload", new Opaque()); + + String guidance = guidance(plugin, args); + + assertThat(guidance).doesNotContain(""); + assertThat(guidance).containsMatch("\"second\" : \\{\\s*\"unit\" : \"F\""); + } + + @Test + public void onToolError_whenANestedArgumentCannotBeSerialized_writesOnlyThatValueAsText() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + ImmutableMap args = + ImmutableMap.of("items", ImmutableList.of(1, new Opaque()), "meta", ImmutableMap.of()); + + String guidance = guidance(plugin, args); + + assertThat(guidance).containsMatch("\"items\" : \\[ 1, \"\" \\]"); + assertThat(guidance).contains("\"meta\" : { }"); + } + + @Test + public void builder_defaults_retryThreeTimesThenLetTheErrorPropagate() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + + callOnError(plugin); + callOnError(plugin); + + assertThat(plugin.getName()).isEqualTo("reflect_retry_tool_plugin"); + assertThat(callOnError(plugin)).containsEntry("retry_count", 3); + plugin + .onToolErrorCallback(tool, ARGS, toolContext, error) + .test() + .assertNoValues() + .assertComplete(); + } + + @Test + public void builder_name_namesThePlugin() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().name("custom").build(); + + assertThat(plugin.getName()).isEqualTo("custom"); + } + + @Test + public void builder_negativeMaxRetries_isRejected() { + assertThrows( + IllegalArgumentException.class, () -> ReflectAndRetryToolPlugin.builder().maxRetries(-1)); + } + + @Test + public void builder_nullTrackingScope_isRejected() { + assertThrows( + NullPointerException.class, () -> ReflectAndRetryToolPlugin.builder().trackingScope(null)); + } + + @Test + public void builder_nullResultErrorExtractor_isRejected() { + assertThrows( + NullPointerException.class, + () -> ReflectAndRetryToolPlugin.builder().resultErrorExtractor(null)); + } + + @Test + public void builder_nullName_isRejected() { + assertThrows(NullPointerException.class, () -> ReflectAndRetryToolPlugin.builder().name(null)); + } + + @Test + public void afterTool_onSuccess_resetsThatToolsCounter() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnError(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterTool_successOfOneTool_doesNotForgiveAnother() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnErrorFor(plugin, otherTool); + + plugin + .afterToolCallback(otherTool, ARGS, toolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + assertThat(callOnErrorFor(plugin, otherTool)).containsEntry("retry_count", 1); + } + + @Test + public void onToolError_countsEachToolSeparately() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnError(plugin); + callOnError(plugin); + + Map other = callOnErrorFor(plugin, otherTool); + + assertThat(other).containsEntry("retry_count", 1); + assertThat(callOnError(plugin)).containsEntry("retry_count", 3); + } + + @Test + public void afterTool_successInOneInvocation_leavesAnotherInvocationsCount() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnErrorIn(plugin, otherToolContext); + callOnError(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnErrorIn(plugin, otherToolContext)).containsEntry("retry_count", 2); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterTool_ownReflectionResponse_isPassedThroughUncounted() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + Map reflection = callOnError(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, reflection) + .test() + .assertNoValues() + .assertComplete(); + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + + plugin + .afterToolCallback(tool, ARGS, toolContext, ImmutableMap.of("ok", true)) + .test() + .assertComplete(); + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void invocationScope_countsAreIsolatedPerInvocation() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + + assertThat(callOnErrorIn(plugin, otherToolContext)).containsEntry("retry_count", 1); + } + + @Test + public void globalScope_countsSurviveAcrossInvocations() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder().trackingScope(TrackingScope.GLOBAL).build(); + callOnError(plugin); + + assertThat(callOnErrorIn(plugin, otherToolContext)).containsEntry("retry_count", 2); + } + + @Test + public void globalScope_successInAnotherInvocation_resetsTheCount() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder().trackingScope(TrackingScope.GLOBAL).build(); + callOnError(plugin); + callOnError(plugin); + + plugin + .afterToolCallback(tool, ARGS, otherToolContext, ImmutableMap.of("ok", true)) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterRun_clearsTheEndedInvocationsCounts() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnError(plugin); + + plugin.afterRunCallback(invocation).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void onRunError_clearsTheEndedInvocationsCounts() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnError(plugin); + + plugin.onRunErrorCallback(invocation, error).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void afterRun_leavesOtherInvocationsCountsAlone() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnErrorIn(plugin, otherToolContext); + + plugin.afterRunCallback(otherInvocation).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + assertThat(callOnErrorIn(plugin, otherToolContext)).containsEntry("retry_count", 1); + } + + @Test + public void globalScope_afterRun_keepsCounts() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder().trackingScope(TrackingScope.GLOBAL).build(); + callOnError(plugin); + + plugin.afterRunCallback(invocation).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + @Test + public void onRunError_leavesOtherInvocationsCountsAlone() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().maxRetries(3).build(); + callOnError(plugin); + callOnErrorIn(plugin, otherToolContext); + + plugin.onRunErrorCallback(otherInvocation, error).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + assertThat(callOnErrorIn(plugin, otherToolContext)).containsEntry("retry_count", 1); + } + + @Test + public void globalScope_onRunError_keepsCounts() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder().trackingScope(TrackingScope.GLOBAL).build(); + callOnError(plugin); + + plugin.onRunErrorCallback(invocation, error).test().assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 2); + } + + @Test + public void defaultExtractor_treatsEveryResultAsSuccess() { + ReflectAndRetryToolPlugin plugin = ReflectAndRetryToolPlugin.builder().build(); + callOnError(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, STATUS_ERROR_RESULT) + .test() + .assertNoValues() + .assertComplete(); + + assertThat(callOnError(plugin)).containsEntry("retry_count", 1); + } + + @Test + public void resultErrorExtractor_throwableError_isReportedByItsType() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsException) + .build(); + + Map response = callAfterTool(plugin); + + assertThat(response).containsEntry("retry_count", 1); + assertThat(response).containsEntry("error_type", "IllegalStateException"); + assertThat(response).containsEntry("error_details", STATUS_ERROR_MESSAGE); + } + + @Test + public void resultErrorExtractor_mapError_isReportedAsToolErrorWithTheMapAsData() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsResult) + .build(); + + Map response = callAfterTool(plugin); + + assertThat(response).containsEntry("error_type", "ToolError"); + assertThat(response).containsEntry("error_details", STATUS_ERROR_RESULT); + assertThat((String) response.get("reflection_guidance")).contains("\"status\" : \"error\""); + } + + @Test + public void resultErrorExtractor_textError_isReportedAsToolErrorText() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsText) + .build(); + + Map response = callAfterTool(plugin); + + assertThat(response).containsEntry("error_type", "ToolError"); + assertThat(response).containsEntry("error_details", STATUS_ERROR_MESSAGE); + assertThat((String) response.get("reflection_guidance")) + .contains("```\n%s\n```".formatted(STATUS_ERROR_MESSAGE)); + } + + @Test + public void resultErrorExtractor_pastRetryLimit_emitsTheThrowableError() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(1) + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsException) + .build(); + callAfterTool(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, STATUS_ERROR_RESULT) + .test() + .assertError( + e -> e instanceof IllegalStateException && STATUS_ERROR_MESSAGE.equals(e.getMessage())); + } + + @Test + public void resultErrorExtractor_pastRetryLimit_wrapsANonThrowableErrorGenerically() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(1) + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsResult) + .build(); + callAfterTool(plugin); + + plugin + .afterToolCallback(tool, ARGS, toolContext, STATUS_ERROR_RESULT) + .test() + .assertError( + e -> + e instanceof IllegalStateException + && GENERIC_EXCEEDED_MESSAGE.equals(e.getMessage())); + } + + @Test + public void resultErrorExtractor_pastRetryLimit_whenNotThrowing_returnsGiveUpMessage() { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .maxRetries(1) + .throwExceptionIfRetryExceeded(false) + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsResult) + .build(); + callAfterTool(plugin); + + Map exceeded = callAfterTool(plugin); + + assertThat(exceeded).containsEntry("error_type", "ToolError"); + assertThat(exceeded).containsEntry("retry_count", 1); + assertThat((String) exceeded.get("reflection_guidance")) + .contains("Do not attempt to use the `flaky_tool` tool again"); + } + + @Test + public void toolReturningSingleError_modelReceivesReflectionGuidance() { + TestLlm llm = createTestLlm(call(TOOL_NAME, FAIL), createTextLlmResponse("done")); + + run(llm, tool, ReflectAndRetryToolPlugin.builder().build()).assertComplete(); + + Map response = functionResponseIn(llm.getRequests().get(1)); + assertThat(response).containsEntry(RESPONSE_TYPE_KEY, RESPONSE_TYPE); + assertThat(response).containsEntry("error_type", "IllegalStateException"); + assertThat(response).containsEntry("error_details", BOOM); + assertThat(response).containsEntry("retry_count", 1); + assertThat(hasGuidance(llm.getRequests().get(1))).isTrue(); + } + + @Test + public void successfulCall_resetsTheCountForTheNextFailure() { + TestLlm llm = + createTestLlm( + call(TOOL_NAME, FAIL), + call(TOOL_NAME, FAIL), + call(TOOL_NAME, SUCCEED), + call(TOOL_NAME, FAIL), + createTextLlmResponse("done")); + + run(llm, tool, ReflectAndRetryToolPlugin.builder().build()).assertComplete(); + + assertThat(functionResponseIn(llm.getRequests().get(2))).containsEntry("retry_count", 2); + assertThat(functionResponseIn(llm.getRequests().get(3))).containsEntry("ok", true); + assertThat(functionResponseIn(llm.getRequests().get(4))).containsEntry("retry_count", 1); + } + + @Test + public void pastRetryLimit_runFailsWithTheToolsOwnError() { + TestLlm llm = + createTestLlm(call(TOOL_NAME, FAIL), call(TOOL_NAME, FAIL), createTextLlmResponse("x")); + + run(llm, tool, ReflectAndRetryToolPlugin.builder().maxRetries(1).build()) + .assertError(e -> e instanceof IllegalStateException && BOOM.equals(e.getMessage())); + } + + @Test + public void functionToolThatThrows_reachesThePluginAsAStatusErrorResult() { + TestLlm llm = + createTestLlm(call(THROWING_METHOD, ImmutableMap.of()), createTextLlmResponse("x")); + BaseTool tool = FunctionTool.create(ReflectAndRetryToolPluginTest.class, THROWING_METHOD); + + run(llm, tool, ReflectAndRetryToolPlugin.builder().build()).assertComplete(); + + Map response = functionResponseIn(llm.getRequests().get(1)); + assertThat(response).doesNotContainKey(RESPONSE_TYPE_KEY); + assertThat(response).containsEntry("status", "error"); + } + + @Test + public void functionToolThatThrows_withAStatusExtractor_modelReceivesReflectionGuidance() { + TestLlm llm = + createTestLlm(call(THROWING_METHOD, ImmutableMap.of()), createTextLlmResponse("x")); + BaseTool tool = FunctionTool.create(ReflectAndRetryToolPluginTest.class, THROWING_METHOD); + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder() + .resultErrorExtractor(ReflectAndRetryToolPluginTest::statusErrorAsResult) + .build(); + + run(llm, tool, plugin).assertComplete(); + + Map response = functionResponseIn(llm.getRequests().get(1)); + assertThat(response).containsEntry(RESPONSE_TYPE_KEY, RESPONSE_TYPE); + assertThat(response).containsEntry("error_type", "ToolError"); + assertThat(response).containsEntry("retry_count", 1); + } + + @Test + public void unknownToolName_neverReachesThePlugin() { + TestLlm llm = createTestLlm(call(UNKNOWN_TOOL, ImmutableMap.of()), createTextLlmResponse("x")); + + run(llm, tool, ReflectAndRetryToolPlugin.builder().build()).assertComplete(); + + assertThat(llm.getRequests().stream().anyMatch(ReflectAndRetryToolPluginTest::hasGuidance)) + .isFalse(); + } + + /** + * Every thread fails the same tool from a common start. A lost update repeats a count and a + * dropped one skips it, so the {@code retry_count}s handed back across all threads are exactly + * {@code 1..n} if and only if every update was atomic. + */ + @Test + public void onToolError_inParallel_countsEachFailureExactlyOnce() throws Exception { + ReflectAndRetryToolPlugin plugin = + ReflectAndRetryToolPlugin.builder().maxRetries(TOTAL_FAILURES).build(); + CyclicBarrier startTogether = new CyclicBarrier(THREADS); + ExecutorService threads = Executors.newFixedThreadPool(THREADS); + + List>> counts = + threads.invokeAll(Collections.nCopies(THREADS, failRepeatedlyAt(plugin, startTogether))); + threads.shutdown(); + + assertThat(threads.awaitTermination(PARALLEL_TIMEOUT_SECONDS, TimeUnit.SECONDS)).isTrue(); + assertThat(resolve(counts)).containsExactlyElementsIn(oneTo(TOTAL_FAILURES)); + } + + private Callable> failRepeatedlyAt( + ReflectAndRetryToolPlugin plugin, CyclicBarrier barrier) { + return () -> failRepeatedlyAfterBarrier(plugin, barrier); + } + + private ImmutableList failRepeatedlyAfterBarrier( + ReflectAndRetryToolPlugin plugin, CyclicBarrier barrier) throws Exception { + barrier.await(PARALLEL_TIMEOUT_SECONDS, TimeUnit.SECONDS); + return IntStream.range(0, FAILURES_PER_THREAD) + .mapToObj(unused -> retryCount(callOnError(plugin))) + .collect(ImmutableList.toImmutableList()); + } + + private static int retryCount(Map response) { + return (Integer) response.get("retry_count"); + } + + 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()); + } + + private Map callOnError(ReflectAndRetryToolPlugin plugin) { + return callOnError(plugin, error); + } + + private Map callOnError(ReflectAndRetryToolPlugin plugin, Throwable thrown) { + return plugin.onToolErrorCallback(tool, ARGS, toolContext, thrown).blockingGet(); + } + + private Map callOnErrorFor(ReflectAndRetryToolPlugin plugin, BaseTool failing) { + return plugin.onToolErrorCallback(failing, ARGS, toolContext, error).blockingGet(); + } + + private Map callOnErrorIn(ReflectAndRetryToolPlugin plugin, ToolContext context) { + return plugin.onToolErrorCallback(tool, ARGS, context, error).blockingGet(); + } + + private static InvocationContext invocation(String invocationId) { + return createInvocationContext(createRootAgent()).toBuilder() + .invocationId(invocationId) + .build(); + } + + private String guidance(ReflectAndRetryToolPlugin plugin, Map args) { + return (String) + plugin + .onToolErrorCallback(tool, args, toolContext, error) + .blockingGet() + .get("reflection_guidance"); + } + + private Map callAfterTool(ReflectAndRetryToolPlugin plugin) { + return plugin.afterToolCallback(tool, ARGS, toolContext, STATUS_ERROR_RESULT).blockingGet(); + } + + /** A {@code FunctionTool} method that throws, which {@code FunctionTool} turns into a result. */ + public static Map alwaysThrows() { + throw new IllegalStateException(BOOM); + } + + private static TestSubscriber run( + TestLlm llm, BaseTool tool, ReflectAndRetryToolPlugin plugin) { + LlmAgent agent = createTestAgentBuilder(llm).tools(ImmutableList.of(tool)).build(); + Runner runner = + Runner.builder() + .app( + App.builder() + .name(APP_NAME) + .rootAgent(agent) + .plugins(ImmutableList.of(plugin)) + .build()) + .build(); + Session session = runner.sessionService().createSession(APP_NAME, USER_ID).blockingGet(); + return runner + .runAsync(USER_ID, session.id(), userMessage()) + .test() + .awaitDone(RUN_TIMEOUT_SECONDS, TimeUnit.SECONDS); + } + + private static Content userMessage() { + return Content.builder().role("user").parts(Part.fromText("go")).build(); + } + + /** A model turn calling {@code toolName}, with an id of its own so responses pair correctly. */ + private static LlmResponse call(String toolName, Map args) { + return createFunctionCallLlmResponse(CALL_ID.formatted(UUID.randomUUID()), toolName, args); + } + + /** The function response carried by the last content of {@code request}. */ + private static Map functionResponseIn(LlmRequest request) { + return functionResponses(Iterables.getLast(request.contents())).stream() + .findFirst() + .orElseThrow(); + } + + private static boolean hasGuidance(LlmRequest request) { + return request.contents().stream() + .flatMap(content -> functionResponses(content).stream()) + .anyMatch(response -> RESPONSE_TYPE.equals(response.get(RESPONSE_TYPE_KEY))); + } + + private static ImmutableList> functionResponses(Content content) { + return content.parts().orElse(ImmutableList.of()).stream() + .flatMap(part -> part.functionResponse().stream()) + .flatMap(response -> response.response().stream()) + .collect(ImmutableList.toImmutableList()); + } + + /** Reports {@code {"status": "error"}} as an exception. */ + private static Optional statusErrorAsException( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + return isStatusError(result) + ? Optional.of(new IllegalStateException(STATUS_ERROR_MESSAGE)) + : Optional.empty(); + } + + /** Reports the result itself as the error, the way upstream's documentation suggests. */ + private static Optional statusErrorAsResult( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + return isStatusError(result) ? Optional.of(result) : Optional.empty(); + } + + /** Reports {@code {"status": "error"}} as plain text. */ + private static Optional statusErrorAsText( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + return isStatusError(result) ? Optional.of(STATUS_ERROR_MESSAGE) : Optional.empty(); + } + + private static boolean isStatusError(Map result) { + return "error".equals(result.get("status")); + } + + /** Fails with {@code Single.error} when called with {@code "fail": true}, else succeeds. */ + private static final class FlakyTool extends BaseTool { + + FlakyTool(String name) { + super(name, "Fails on request"); + } + + @Override + public Optional declaration() { + return Optional.of(FunctionDeclaration.builder().name(name()).build()); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + return Boolean.TRUE.equals(args.get("fail")) + ? Single.error(new IllegalStateException(BOOM)) + : Single.just(ImmutableMap.of("ok", true)); + } + } + + /** 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 ""; + } + } +}