Unwrap and rethrow specificied exceptions DefaultToolExecutionExceptionProcessor

Signed-off-by: Daniel Garnier-Moiroux <[email protected]>

Author: Kehrlann

spring-ai-model/src/main/java/org/springframework/ai/tool/execution/DefaultToolExecutionExceptionProcessor.java

@@ -16,15 +16,20 @@
 
 package org.springframework.ai.tool.execution;
 
+import java.util.Collections;
+import java.util.List;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import org.springframework.util.Assert;
 
 /**
- * Default implementation of {@link ToolExecutionExceptionProcessor}.
+ * Default implementation of {@link ToolExecutionExceptionProcessor}. Can be configured
+ * with an allowlist of exceptions that will be unwrapped from the
+ * {@link ToolExecutionException} and rethrown as is.
  *
  * @author Thomas Vitale
+ * @author Daniel Garnier-Moiroux
  * @since 1.0.0
  */
 public class DefaultToolExecutionExceptionProcessor implements ToolExecutionExceptionProcessor {
@@ -35,13 +40,26 @@ public class DefaultToolExecutionExceptionProcessor implements ToolExecutionExce
 
 	private final boolean alwaysThrow;
 
+	private final List<Class<? extends RuntimeException>> rethrownExceptions;
+
 	public DefaultToolExecutionExceptionProcessor(boolean alwaysThrow) {
+		this(alwaysThrow, Collections.emptyList());
+	}
+
+	public DefaultToolExecutionExceptionProcessor(boolean alwaysThrow,
+			List<Class<? extends RuntimeException>> rethrownExceptions) {
 		this.alwaysThrow = alwaysThrow;
+		this.rethrownExceptions = Collections.unmodifiableList(rethrownExceptions);
 	}
 
 	@Override
 	public String process(ToolExecutionException exception) {
 		Assert.notNull(exception, "exception cannot be null");
+		Throwable cause = exception.getCause();
+		if (cause instanceof RuntimeException runtimeException
+				&& this.rethrownExceptions.stream().anyMatch(rethrown -> rethrown.isAssignableFrom(cause.getClass()))) {
+			throw runtimeException;
+		}
 		if (this.alwaysThrow) {
 			throw exception;
 		}
@@ -58,13 +76,31 @@ public static class Builder {
 
 		private boolean alwaysThrow = DEFAULT_ALWAYS_THROW;
 
+		private List<Class<? extends RuntimeException>> exceptions = Collections.emptyList();
+
+		/**
+		 * Rethrow the {@link ToolExecutionException}
+		 * @param alwaysThrow when true, throws; when false, returns the exception message
+		 * @return the builder instance
+		 */
 		public Builder alwaysThrow(boolean alwaysThrow) {
 			this.alwaysThrow = alwaysThrow;
 			return this;
 		}
 
+		/**
+		 * An allowlist of exceptions thrown by tools, which will be unwrapped and
+		 * re-thrown without further processing.
+		 * @param exceptions the list of exceptions
+		 * @return the builder instance
+		 */
+		public Builder rethrowExceptions(List<Class<? extends RuntimeException>> exceptions) {
+			this.exceptions = exceptions;
+			return this;
+		}
+
 		public DefaultToolExecutionExceptionProcessor build() {
-			return new DefaultToolExecutionExceptionProcessor(this.alwaysThrow);
+			return new DefaultToolExecutionExceptionProcessor(this.alwaysThrow, exceptions);
 		}
 
 	}

spring-ai-model/src/test/java/org/springframework/ai/tool/execution/DefaultToolExecutionExceptionProcessorTests.java

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2023-2025 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 org.springframework.ai.tool.execution;
+
+import java.util.List;
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.tool.definition.DefaultToolDefinition;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.InstanceOfAssertFactories.type;
+
+/**
+ * Unit tests for {@link DefaultToolExecutionExceptionProcessor}.
+ *
+ * @author Daniel Garnier-Moiroux
+ */
+class DefaultToolExecutionExceptionProcessorTests {
+
+	private final IllegalStateException toolException = new IllegalStateException("Inner exception");
+
+	private final DefaultToolDefinition toolDefinition = new DefaultToolDefinition("toolName", "toolDescription",
+			"inputSchema");
+
+	private final ToolExecutionException toolExecutionException = new ToolExecutionException(toolDefinition,
+			toolException);
+
+	@Test
+	void processReturnsMessage() {
+		DefaultToolExecutionExceptionProcessor processor = DefaultToolExecutionExceptionProcessor.builder().build();
+
+		String result = processor.process(this.toolExecutionException);
+
+		assertThat(result).isEqualTo(this.toolException.getMessage());
+	}
+
+	@Test
+	void processAlwaysThrows() {
+		DefaultToolExecutionExceptionProcessor processor = DefaultToolExecutionExceptionProcessor.builder()
+			.alwaysThrow(true)
+			.build();
+
+		assertThatThrownBy(() -> processor.process(this.toolExecutionException))
+			.hasMessage(this.toolException.getMessage())
+			.hasCauseInstanceOf(this.toolException.getClass())
+			.asInstanceOf(type(ToolExecutionException.class))
+			.extracting(ToolExecutionException::getToolDefinition)
+			.isEqualTo(this.toolDefinition);
+	}
+
+	@Test
+	void processRethrows() {
+		DefaultToolExecutionExceptionProcessor processor = DefaultToolExecutionExceptionProcessor.builder()
+			.alwaysThrow(false)
+			.rethrowExceptions(List.of(IllegalStateException.class))
+			.build();
+
+		assertThatThrownBy(() -> processor.process(this.toolExecutionException)).isEqualTo(this.toolException);
+	}
+
+	@Test
+	void processRethrowsExceptionSubclasses() {
+		DefaultToolExecutionExceptionProcessor processor = DefaultToolExecutionExceptionProcessor.builder()
+			.alwaysThrow(false)
+			.rethrowExceptions(List.of(RuntimeException.class))
+			.build();
+
+		assertThatThrownBy(() -> processor.process(this.toolExecutionException)).isEqualTo(this.toolException);
+	}
+
+	@Test
+	void processRethrowsOnlySelectExceptions() {
+		DefaultToolExecutionExceptionProcessor processor = DefaultToolExecutionExceptionProcessor.builder()
+			.alwaysThrow(false)
+			.rethrowExceptions(List.of(IllegalStateException.class))
+			.build();
+
+		ToolExecutionException exception = new ToolExecutionException(this.toolDefinition,
+				new RuntimeException("This exception was not rethrown"));
+		String result = processor.process(exception);
+
+		assertThat(result).isEqualTo("This exception was not rethrown");
+	}
+
+}