ToolContext is now excluded from JSON Schema Generation

When generating the JSON Schema for a tool input from a method, ToolContext is now excluded since it's not something we want the model to provide. The framework takes care of passing a value for it when actually executing the tool call.

Fixes gh-2366

Signed-off-by: Thomas Vitale <[email protected]>

Author: ThomasVitale

spring-ai-core/src/main/java/org/springframework/ai/chat/model/ToolContext.java

@@ -43,7 +43,7 @@
  * @author Christian Tzolov
  * @since 1.0.0
  */
-public class ToolContext {
+public final class ToolContext {
 
 	/**
 	 * The key for the running, tool call history stored in the context map.

spring-ai-core/src/main/java/org/springframework/ai/util/json/schema/JsonSchemaGenerator.java

@@ -31,10 +31,12 @@
 import com.github.victools.jsonschema.module.jackson.JacksonOption;
 import com.github.victools.jsonschema.module.swagger2.Swagger2Module;
 import io.swagger.v3.oas.annotations.media.Schema;
+import org.springframework.ai.chat.model.ToolContext;
 import org.springframework.ai.tool.annotation.ToolParam;
 import org.springframework.ai.util.json.JsonParser;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
 import org.springframework.util.StringUtils;
 
 import java.lang.reflect.Method;
@@ -125,6 +127,14 @@ public static String generateForMethodInput(Method method, SchemaOption... schem
 		for (int i = 0; i < method.getParameterCount(); i++) {
 			String parameterName = method.getParameters()[i].getName();
 			Type parameterType = method.getGenericParameterTypes()[i];
+			if (parameterType instanceof Class<?> parameterClass
+					&& ClassUtils.isAssignable(parameterClass, ToolContext.class)) {
+				// A ToolContext method parameter is not included in the JSON Schema
+				// generation.
+				// It's a special type used by Spring AI to pass contextual data to tools
+				// outside the model interaction flow.
+				continue;
+			}
 			if (isMethodParameterRequired(method, i)) {
 				required.add(parameterName);
 			}

spring-ai-core/src/test/java/org/springframework/ai/util/json/JsonSchemaGeneratorTests.java

@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.databind.JsonNode;
 import io.swagger.v3.oas.annotations.media.Schema;
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ToolContext;
 import org.springframework.ai.tool.annotation.ToolParam;
 import org.springframework.ai.util.json.schema.JsonSchemaGenerator;
 import org.springframework.lang.Nullable;
@@ -347,6 +348,36 @@ void generateSchemaForMethodWithTimeParameters() throws Exception {
 		assertThat(schema).isEqualToIgnoringWhitespace(expectedJsonSchema);
 	}
 
+	@Test
+	void generateSchemaForMethodWithToolContext() throws Exception {
+		Method method = TestMethods.class.getDeclaredMethod("contextMethod", String.class, LocalDateTime.class,
+				ToolContext.class);
+
+		String schema = JsonSchemaGenerator.generateForMethodInput(method);
+		String expectedJsonSchema = """
+				{
+				    "$schema": "https://json-schema.org/draft/2020-12/schema",
+				    "type": "object",
+				    "properties": {
+				        "deliveryStatus": {
+				            "type": "string"
+				        },
+				        "expectedDelivery": {
+				            "type": "string",
+				            "format": "date-time"
+				        }
+				    },
+				    "required": [
+				        "deliveryStatus",
+				        "expectedDelivery"
+				    ],
+				    "additionalProperties": false
+				}
+				""";
+
+		assertThat(schema).isEqualToIgnoringWhitespace(expectedJsonSchema);
+	}
+
 	// TYPES
 
 	@Test
@@ -658,6 +689,9 @@ public void complexMethod(List<String> items, TestData data, MoreTestData moreDa
 		public void timeMethod(Duration duration, LocalDateTime localDateTime, Instant instant) {
 		}
 
+		public void contextMethod(String deliveryStatus, LocalDateTime expectedDelivery, ToolContext toolContext) {
+		}
+
 	}
 
 	record TestData(int id, @ToolParam(description = "The special name") String name) {

spring-ai-docs/src/main/antora/modules/ROOT/pages/upgrade-notes.adoc

@@ -171,9 +171,11 @@ To use this automation:
 
 This approach can save time and reduce the chance of errors when upgrading multiple projects or complex codebases.
 
+== Upgrading to 1.0.0.M7
 
-== Upgrading to 1.0.0.M6
+* The `ToolContext` class has now been marked as final and cannot be extended anymore. It was never supposed to be subclassed.
 
+== Upgrading to 1.0.0.M6
 
 === Changes to Usage Interface and DefaultUsage Implementation