fix(agent): Improving the retry semantics of ToolRetryInterceptor (#4794)

Duansg · web-flow · commit 102336f840d3 · 2026-07-09T23:29:09.000+08:00
diff --git a/spring-ai-alibaba-agent-framework/src/main/java/com/alibaba/cloud/ai/graph/agent/interceptor/toolretry/ToolRetryInterceptor.java b/spring-ai-alibaba-agent-framework/src/main/java/com/alibaba/cloud/ai/graph/agent/interceptor/toolretry/ToolRetryInterceptor.java
@@ -45,7 +45,7 @@ public class ToolRetryInterceptor extends ToolInterceptor {
 
 	private static final Logger log = LoggerFactory.getLogger(ToolRetryInterceptor.class);
 
-	private final int maxRetries;
+	private final int maxAttempts;
 	private final Set<String> toolNames;
 	private final Predicate<Exception> retryOn;
 	private final OnFailureBehavior onFailure;
@@ -56,7 +56,7 @@ public class ToolRetryInterceptor extends ToolInterceptor {
 	private final boolean jitter;
 
 	private ToolRetryInterceptor(Builder builder) {
-		this.maxRetries = builder.maxRetries;
+		this.maxAttempts = builder.maxAttempts;
 		this.toolNames = builder.toolNames != null ? new HashSet<>(builder.toolNames) : null;
 		this.retryOn = builder.retryOn;
 		this.onFailure = builder.onFailure;
@@ -81,18 +81,17 @@ public ToolCallResponse interceptToolCall(ToolCallRequest request, ToolCallHandl
 		}
 
 		Exception lastException = null;
-		int attempt = 0;
 
-		while (attempt < maxRetries) {
+		// maxAttempts counts the initial call plus retries and is always >= 1, so the tool is executed at least once and is never skipped.
+		for (int attempt = 0; attempt < maxAttempts; attempt++) {
 			try {
-                ToolCallResponse response = handler.call(request);
-                if (ToolCallResponse.SUCCESS_STATUS.equals(response.getStatus())) {
-                    return response;
-                }
-                else {
-                    //  fail, continue run in  while method
-                    throw new RuntimeException(response.getResult().toString());
-                }
+				ToolCallResponse response = handler.call(request);
+				if (ToolCallResponse.SUCCESS_STATUS.equals(response.getStatus())) {
+					return response;
+				}
+				// A non-success status is treated as a failure so it can be retried.
+				String result = response.getResult() != null ? response.getResult() : "unknown error";
+				throw new RuntimeException(result);
 			}
 			catch (Exception e) {
 				lastException = e;
@@ -103,15 +102,15 @@ public ToolCallResponse interceptToolCall(ToolCallRequest request, ToolCallHandl
 					throw e;
 				}
 
-				if (attempt == maxRetries) {
-					// Max retries reached
+				// Last attempt failed: Stop the operation
+				if (attempt >= maxAttempts - 1) {
 					break;
 				}
 
 				// Calculate delay
 				long delay = calculateDelay(attempt);
 				log.warn("Tool '{}' failed (attempt {}/{}), retrying in {}ms: {}",
-						toolName, attempt + 1, maxRetries + 1, delay, e.getMessage());
+						toolName, attempt + 1, maxAttempts, delay, e.getMessage());
 
 				try {
 					Thread.sleep(delay);
@@ -120,22 +119,20 @@ public ToolCallResponse interceptToolCall(ToolCallRequest request, ToolCallHandl
 					Thread.currentThread().interrupt();
 					throw new RuntimeException("Retry interrupted", ie);
 				}
-
-				attempt++;
 			}
 		}
 
 		// All retries exhausted
 		if (onFailure == OnFailureBehavior.RAISE) {
-			throw new RuntimeException("Tool call failed after " + (maxRetries + 1) + " attempts", lastException);
+			throw new RuntimeException("Tool call failed after " + maxAttempts + " attempts", lastException);
 		}
 		else {
 			// Return error message as tool response
 			String errorMessage = errorFormatter != null
 					? errorFormatter.apply(lastException)
-					: "Tool call failed after " + (maxRetries + 1) + " attempts: " + lastException.getMessage();
+					: "Tool call failed after " + maxAttempts + " attempts: " + lastException.getMessage();
 
-			log.error("Tool '{}' failed after {} attempts: {}", toolName, maxRetries + 1, lastException.getMessage());
+			log.error("Tool '{}' failed after {} attempts: {}", toolName, maxAttempts, lastException.getMessage());
 			return ToolCallResponse.of(request.getToolCallId(), request.getToolName(), errorMessage);
 		}
 	}
@@ -164,7 +161,8 @@ public enum OnFailureBehavior {
 	}
 
 	public static class Builder {
-		private int maxRetries = 2;
+		// Total number of attempts including the first call.
+		private int maxAttempts = 3;
 		private Set<String> toolNames;
 		private Predicate<Exception> retryOn = e -> true; // Retry on all exceptions by default
 		private OnFailureBehavior onFailure = OnFailureBehavior.RETURN_MESSAGE;
@@ -174,11 +172,29 @@ public static class Builder {
 		private long maxDelayMs = 60000;
 		private boolean jitter = true;
 
+		/**
+		 * Set the maximum number of attempts, including the first call. Must be >= 1.
+		 * @param maxAttempts total attempts (initial call plus retries)
+		 */
+		public Builder maxAttempts(int maxAttempts) {
+			if (maxAttempts < 1) {
+				throw new IllegalArgumentException("maxAttempts must be >= 1");
+			}
+			this.maxAttempts = maxAttempts;
+			return this;
+		}
+
+		/**
+		 * Set the maximum number of retries (excluding the first call).
+		 * @param maxRetries number of retries; total attempts will be {@code maxRetries + 1}
+		 * @deprecated use {@link #maxAttempts(int)} instead, which counts the initial call
+		 */
+		@Deprecated
 		public Builder maxRetries(int maxRetries) {
 			if (maxRetries < 0) {
 				throw new IllegalArgumentException("maxRetries must be >= 0");
 			}
-			this.maxRetries = maxRetries;
+			this.maxAttempts = maxRetries + 1;
 			return this;
 		}
 
diff --git a/spring-ai-alibaba-agent-framework/src/test/java/com/alibaba/cloud/ai/graph/agent/interceptors/ToolRetryInterceptorUnitTest.java b/spring-ai-alibaba-agent-framework/src/test/java/com/alibaba/cloud/ai/graph/agent/interceptors/ToolRetryInterceptorUnitTest.java
@@ -0,0 +1,146 @@
+/*
+ * Copyright 2024-2026 the original author or authors.
+ *
+ * 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
+ *
+ *      https://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.alibaba.cloud.ai.graph.agent.interceptors;
+
+import com.alibaba.cloud.ai.graph.agent.interceptor.ToolCallRequest;
+import com.alibaba.cloud.ai.graph.agent.interceptor.ToolCallResponse;
+import com.alibaba.cloud.ai.graph.agent.interceptor.toolretry.ToolRetryInterceptor;
+
+import java.util.HashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * Pure unit tests for {@link ToolRetryInterceptor} that exercise the retry loop directly
+ * with a mock {@link com.alibaba.cloud.ai.graph.agent.interceptor.ToolCallHandler}, so no
+ * model / API key is required. Guards against the off-by-one regression where the loop
+ * performed one attempt too few, skipped the tool entirely for a single-attempt config
+ * (leading to an NPE), and slept an unused backoff after the final failure. Also verifies
+ * the {@code maxAttempts} API and the deprecated {@code maxRetries} alias stay in sync.
+ */
+class ToolRetryInterceptorUnitTest {
+
+	private ToolCallRequest request() {
+		return new ToolCallRequest("t", "{}", "id-1", new HashMap<>());
+	}
+
+	private ToolRetryInterceptor.Builder builder(int maxAttempts) {
+		return ToolRetryInterceptor.builder().maxAttempts(maxAttempts).initialDelay(1).jitter(false);
+	}
+
+	@Test
+	void runsAllConfiguredAttempts() {
+		AtomicInteger calls = new AtomicInteger();
+		ToolCallResponse response = builder(3).build().interceptToolCall(request(), r -> {
+			calls.incrementAndGet();
+			throw new RuntimeException("boom");
+		});
+		assertEquals(3, calls.get());
+		assertTrue(response.getResult().contains("3 attempts"));
+	}
+
+	@Test
+	void twoAttemptsMeansOneRetry() {
+		AtomicInteger calls = new AtomicInteger();
+		builder(2).build().interceptToolCall(request(), r -> {
+			calls.incrementAndGet();
+			throw new RuntimeException("boom");
+		});
+		assertEquals(2, calls.get());
+	}
+
+	@Test
+	void singleAttemptRunsToolExactlyOnceWithoutNpe() {
+		AtomicInteger calls = new AtomicInteger();
+		ToolCallResponse response = builder(1).build().interceptToolCall(request(), r -> {
+			calls.incrementAndGet();
+			throw new RuntimeException("boom");
+		});
+		// The tool must still execute once, and no NullPointerException must be thrown.
+		assertEquals(1, calls.get());
+		assertTrue(response.getResult().contains("1 attempts"));
+	}
+
+	@Test
+	void stopsRetryingOnceSuccessful() {
+		AtomicInteger calls = new AtomicInteger();
+		ToolCallResponse response = builder(6).build().interceptToolCall(request(), r -> {
+			if (calls.incrementAndGet() < 3) {
+				throw new RuntimeException("transient");
+			}
+			return ToolCallResponse.success(r.getToolCallId(), r.getToolName(), "ok");
+		});
+		assertEquals(3, calls.get());
+		assertEquals(ToolCallResponse.SUCCESS_STATUS, response.getStatus());
+		assertEquals("ok", response.getResult());
+	}
+
+	@Test
+	void raiseBehaviorThrowsAfterAllAttempts() {
+		AtomicInteger calls = new AtomicInteger();
+		ToolRetryInterceptor interceptor = builder(3)
+			.onFailure(ToolRetryInterceptor.OnFailureBehavior.RAISE)
+			.build();
+		RuntimeException ex = assertThrows(RuntimeException.class,
+				() -> interceptor.interceptToolCall(request(), r -> {
+					calls.incrementAndGet();
+					throw new RuntimeException("boom");
+				}));
+		assertEquals(3, calls.get());
+		assertTrue(ex.getMessage().contains("3 attempts"));
+	}
+
+	@Test
+	void nonRetryableExceptionIsRethrownImmediately() {
+		AtomicInteger calls = new AtomicInteger();
+		ToolRetryInterceptor interceptor = builder(4)
+			.retryOn(IllegalStateException.class) // only retry ISE
+			.build();
+		assertThrows(IllegalArgumentException.class,
+				() -> interceptor.interceptToolCall(request(), r -> {
+					calls.incrementAndGet();
+					throw new IllegalArgumentException("not retryable");
+				}));
+		assertEquals(1, calls.get());
+	}
+
+	@Test
+	@SuppressWarnings("deprecation")
+	void deprecatedMaxRetriesAliasMapsToAttempts() {
+		// maxRetries(n) == maxAttempts(n + 1)
+		AtomicInteger calls = new AtomicInteger();
+		ToolRetryInterceptor.builder().maxRetries(2).initialDelay(1).jitter(false).build()
+			.interceptToolCall(request(), r -> {
+				calls.incrementAndGet();
+				throw new RuntimeException("boom");
+			});
+		assertEquals(3, calls.get());
+
+		// The previously buggy edge case: maxRetries(0) must still run the tool once, not zero times.
+		AtomicInteger zeroCalls = new AtomicInteger();
+		ToolRetryInterceptor.builder().maxRetries(0).initialDelay(1).jitter(false).build()
+			.interceptToolCall(request(), r -> {
+				zeroCalls.incrementAndGet();
+				throw new RuntimeException("boom");
+			});
+		assertEquals(1, zeroCalls.get());
+	}
+}
