Wrap Exception from toolFunction in ToolExecutionException and rethrow (#3918)

Fixes #2857 - When `toolFunction` throws a `ToolExecutionException`, rethrow it directly. - When `toolFunction` throws a `Exception`, wrap it in a `ToolExecutionException` and rethrow it. - Add related tests Signed-off-by: YunKui Lu <luyunkui95@gmail.com> (cherry picked from commit 1b5cd09)

Author: YunKuiLu

spring-ai-model/src/main/java/org/springframework/ai/tool/function/FunctionToolCallback.java

@@ -31,6 +31,7 @@
 import org.springframework.ai.tool.definition.ToolDefinition;
 import org.springframework.ai.tool.execution.DefaultToolCallResultConverter;
 import org.springframework.ai.tool.execution.ToolCallResultConverter;
+import org.springframework.ai.tool.execution.ToolExecutionException;
 import org.springframework.ai.tool.metadata.ToolMetadata;
 import org.springframework.ai.tool.support.ToolUtils;
 import org.springframework.ai.util.json.JsonParser;
@@ -44,6 +45,7 @@
  * A {@link ToolCallback} implementation to invoke functions as tools.
  *
  * @author Thomas Vitale
+ * @author YunKui Lu
  * @since 1.0.0
  */
 public class FunctionToolCallback<I, O> implements ToolCallback {
@@ -99,13 +101,25 @@ public String call(String toolInput, @Nullable ToolContext toolContext) {
 logger.debug("Starting execution of tool: {}", this.toolDefinition.name());
 
 I request = JsonParser.fromJson(toolInput, this.toolInputType);
-O response = this.toolFunction.apply(request, toolContext);
+O response = callMethod(request, toolContext);
 
 logger.debug("Successful execution of tool: {}", this.toolDefinition.name());
 
 return this.toolCallResultConverter.convert(response, null);
 }
 
+private O callMethod(I request, @Nullable ToolContext toolContext) {
+try {
+return this.toolFunction.apply(request, toolContext);
+}
+catch (ToolExecutionException ex) {
+throw ex;
+}
+catch (Exception ex) {
+throw new ToolExecutionException(this.toolDefinition, ex);
+}
+}
+
 @Override
 public String toString() {
 return "FunctionToolCallback{" + "toolDefinition=" + this.toolDefinition + ", toolMetadata=" + this.toolMetadata

spring-ai-model/src/test/java/org/springframework/ai/tool/function/FunctionToolCallbackTest.java

@@ -0,0 +1,187 @@
+/*
+ * Copyright 2025-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.function;
+
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiFunction;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.ai.tool.execution.ToolExecutionException;
+import org.springframework.ai.tool.metadata.ToolMetadata;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.assertj.core.api.InstanceOfAssertFactories.type;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * @author YunKui Lu
+ */
+class FunctionToolCallbackTest {
+
+@Test
+void testConsumerToolCall() {
+TestFunctionTool tool = new TestFunctionTool();
+FunctionToolCallback<String, Void> callback = FunctionToolCallback.builder("testTool", tool.stringConsumer())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(String.class)
+.build();
+
+callback.call("\"test string param\"");
+
+assertEquals("test string param", tool.calledValue.get());
+}
+
+@Test
+void testBiFunctionToolCall() {
+TestFunctionTool tool = new TestFunctionTool();
+FunctionToolCallback<String, String> callback = FunctionToolCallback
+.builder("testTool", tool.stringBiFunction())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(String.class)
+.build();
+
+ToolContext toolContext = new ToolContext(Map.of("foo", "bar"));
+
+String callResult = callback.call("\"test string param\"", toolContext);
+
+assertEquals("test string param", tool.calledValue.get());
+assertEquals("\"return value = test string param\"", callResult);
+assertEquals(toolContext, tool.calledToolContext.get());
+}
+
+@Test
+void testFunctionToolCall() {
+TestFunctionTool tool = new TestFunctionTool();
+FunctionToolCallback<String, String> callback = FunctionToolCallback.builder("testTool", tool.stringFunction())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(String.class)
+.build();
+
+ToolContext toolContext = new ToolContext(Map.of());
+
+String callResult = callback.call("\"test string param\"", toolContext);
+
+assertEquals("test string param", tool.calledValue.get());
+assertEquals("\"return value = test string param\"", callResult);
+}
+
+@Test
+void testSupplierToolCall() {
+TestFunctionTool tool = new TestFunctionTool();
+
+FunctionToolCallback<Void, String> callback = FunctionToolCallback.builder("testTool", tool.stringSupplier())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(Void.class)
+.build();
+
+ToolContext toolContext = new ToolContext(Map.of());
+
+String callResult = callback.call("\"test string param\"", toolContext);
+
+assertEquals("not params", tool.calledValue.get());
+assertEquals("\"return value = \"", callResult);
+}
+
+@Test
+void testThrowRuntimeException() {
+TestFunctionTool tool = new TestFunctionTool();
+FunctionToolCallback<String, Void> callback = FunctionToolCallback
+.builder("testTool", tool.throwRuntimeException())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(String.class)
+.build();
+
+assertThatThrownBy(() -> callback.call("\"test string param\"")).hasMessage("test exception")
+.hasCauseInstanceOf(RuntimeException.class)
+.asInstanceOf(type(ToolExecutionException.class))
+.extracting(ToolExecutionException::getToolDefinition)
+.isEqualTo(callback.getToolDefinition());
+}
+
+@Test
+void testThrowToolExecutionException() {
+TestFunctionTool tool = new TestFunctionTool();
+FunctionToolCallback<String, Void> callback = FunctionToolCallback
+.builder("testTool", tool.throwToolExecutionException())
+.toolMetadata(ToolMetadata.builder().returnDirect(true).build())
+.description("test description")
+.inputType(String.class)
+.build();
+
+assertThatThrownBy(() -> callback.call("\"test string param\"")).hasMessage("test exception")
+.hasCauseInstanceOf(RuntimeException.class)
+.isInstanceOf(ToolExecutionException.class);
+}
+
+static class TestFunctionTool {
+
+AtomicReference<Object> calledValue = new AtomicReference<>();
+
+AtomicReference<ToolContext> calledToolContext = new AtomicReference<>();
+
+public Consumer<String> stringConsumer() {
+return s -> {
+calledValue.set(s);
+};
+}
+
+public BiFunction<String, ToolContext, String> stringBiFunction() {
+return (s, context) -> {
+calledValue.set(s);
+calledToolContext.set(context);
+return "return value = " + s;
+};
+}
+
+public Function<String, String> stringFunction() {
+return s -> {
+calledValue.set(s);
+return "return value = " + s;
+};
+}
+
+public Supplier<String> stringSupplier() {
+calledValue.set("not params");
+return () -> "return value = ";
+}
+
+public Consumer<String> throwRuntimeException() {
+return s -> {
+throw new RuntimeException("test exception");
+};
+}
+
+public Consumer<String> throwToolExecutionException() {
+return s -> {
+throw new ToolExecutionException(null, new RuntimeException("test exception"));
+};
+}
+
+}
+
+}