diff --git a/spring-ai-model/src/main/java/org/springframework/ai/model/tool/DefaultToolCallingManager.java b/spring-ai-model/src/main/java/org/springframework/ai/model/tool/DefaultToolCallingManager.java index 5149a98a85c..17c6b5b987b 100644 --- a/spring-ai-model/src/main/java/org/springframework/ai/model/tool/DefaultToolCallingManager.java +++ b/spring-ai-model/src/main/java/org/springframework/ai/model/tool/DefaultToolCallingManager.java @@ -34,6 +34,7 @@ import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.consent.ConsentAwareToolCallback; import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.ai.tool.execution.DefaultToolExecutionExceptionProcessor; import org.springframework.ai.tool.execution.ToolExecutionException; @@ -217,7 +218,12 @@ private InternalToolExecutionResult executeToolCall(Prompt prompt, AssistantMess .observe(() -> { String toolResult; try { - toolResult = toolCallback.call(toolInputArguments, toolContext); + if (toolCallback instanceof ConsentAwareToolCallback consentAwareCallback) { + toolResult = consentAwareCallback.call(toolInputArguments, toolContext); + } + else { + toolResult = toolCallback.call(toolInputArguments, toolContext); + } } catch (ToolExecutionException ex) { toolResult = this.toolExecutionExceptionProcessor.process(ex); diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/annotation/RequiresConsent.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/annotation/RequiresConsent.java new file mode 100644 index 00000000000..4368d3c4d88 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/annotation/RequiresConsent.java @@ -0,0 +1,92 @@ +/* + * 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.annotation; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Annotation to indicate that a tool method requires user consent before execution. When + * applied to a method annotated with {@link Tool}, the execution will be intercepted to + * request user approval before proceeding. + * + *

+ * Example usage: + * 

{@code
+ * @Tool(description = "Deletes a book from the database")
+ * @RequiresConsent(message = "The book {bookId} will be permanently deleted. Do you approve?")
+ * public void deleteBook(String bookId) {
+ *     // Implementation
+ * }
+ * }
+ * + * @author Hyunjoon Park + * @since 1.0.0 + */ +@Target(ElementType.METHOD) +@Retention(RetentionPolicy.RUNTIME) +@Documented +public @interface RequiresConsent { + + /** + * The message to display when requesting consent. Supports placeholder syntax using + * curly braces (e.g., {paramName}) which will be replaced with actual parameter + * values at runtime. + * @return the consent message template + */ + String message() default "This action requires your approval. Do you want to proceed?"; + + /** + * The level of consent required. This can be used to implement different consent + * strategies (e.g., one-time consent, session-based consent, etc.). + * @return the consent level + */ + ConsentLevel level() default ConsentLevel.EVERY_TIME; + + /** + * Optional categories for grouping consent requests. This can be used to manage + * consent preferences by category. + * @return array of consent categories + */ + String[] categories() default {}; + + /** + * Defines the consent level for tool execution. + */ + enum ConsentLevel { + + /** + * Requires consent every time the tool is called. + */ + EVERY_TIME, + + /** + * Requires consent once per session. + */ + SESSION, + + /** + * Requires consent once and remembers the preference. + */ + REMEMBER + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareMethodToolCallbackProvider.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareMethodToolCallbackProvider.java new file mode 100644 index 00000000000..a1357a5e1d9 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareMethodToolCallbackProvider.java @@ -0,0 +1,143 @@ +/* + * 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.consent; + +import java.lang.reflect.Method; +import java.util.Arrays; +import java.util.List; + +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.method.MethodToolCallbackProvider; +import org.springframework.core.annotation.AnnotationUtils; +import org.springframework.util.Assert; + +/** + * Extension of {@link MethodToolCallbackProvider} that wraps tool callbacks requiring + * consent with {@link ConsentAwareToolCallback}. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +public class ConsentAwareMethodToolCallbackProvider extends MethodToolCallbackProvider { + + private final ConsentManager consentManager; + + /** + * Creates a new consent-aware method tool callback provider. + * @param toolObjects the objects containing tool methods + * @param consentManager the consent manager for handling consent requests + */ + public ConsentAwareMethodToolCallbackProvider(List toolObjects, ConsentManager consentManager) { + super(toolObjects); + Assert.notNull(consentManager, "consentManager must not be null"); + this.consentManager = consentManager; + } + + @Override + public ToolCallback[] getToolCallbacks() { + ToolCallback[] callbacks = super.getToolCallbacks(); + + // Wrap callbacks that require consent + for (int i = 0; i < callbacks.length; i++) { + ToolCallback callback = callbacks[i]; + RequiresConsent requiresConsent = findRequiresConsentAnnotation(callback); + + if (requiresConsent != null) { + callbacks[i] = new ConsentAwareToolCallback(callback, this.consentManager, requiresConsent); + } + } + + return callbacks; + } + + /** + * Finds the @RequiresConsent annotation for a tool callback. This method checks the + * original method that the callback was created from. + * @param callback the tool callback + * @return the RequiresConsent annotation or null if not present + */ + private RequiresConsent findRequiresConsentAnnotation(ToolCallback callback) { + // For MethodToolCallback, we need to find the original method + // This requires accessing the method through reflection or storing it + // For now, we'll check all methods in the tool objects + + for (Object toolObject : getToolObjects()) { + Method[] methods = toolObject.getClass().getDeclaredMethods(); + for (Method method : methods) { + // Check if this method corresponds to the callback + if (method.getName().equals(callback.getName())) { + RequiresConsent annotation = AnnotationUtils.findAnnotation(method, RequiresConsent.class); + if (annotation != null) { + return annotation; + } + } + } + } + + return null; + } + + /** + * Gets the list of tool objects from the parent class. This is a workaround since the + * field is private in the parent. + * @return the list of tool objects + */ + private List getToolObjects() { + // This would need to be implemented properly, possibly by: + // 1. Making the field protected in the parent class + // 2. Adding a getter in the parent class + // 3. Storing a copy in this class + // For now, we'll throw an exception indicating this needs to be addressed + throw new UnsupportedOperationException( + "Need to access tool objects from parent class. Consider making the field protected or adding a getter."); + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + + private List toolObjects; + + private ConsentManager consentManager; + + private Builder() { + } + + public Builder toolObjects(Object... toolObjects) { + Assert.notNull(toolObjects, "toolObjects cannot be null"); + this.toolObjects = Arrays.asList(toolObjects); + return this; + } + + public Builder consentManager(ConsentManager consentManager) { + Assert.notNull(consentManager, "consentManager cannot be null"); + this.consentManager = consentManager; + return this; + } + + public ConsentAwareMethodToolCallbackProvider build() { + Assert.notNull(this.toolObjects, "toolObjects must be set"); + Assert.notNull(this.consentManager, "consentManager must be set"); + return new ConsentAwareMethodToolCallbackProvider(this.toolObjects, this.consentManager); + } + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareToolCallback.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareToolCallback.java new file mode 100644 index 00000000000..3a73584efcc --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentAwareToolCallback.java @@ -0,0 +1,154 @@ +/* + * 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.consent; + +import java.util.Map; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; + +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.consent.exception.ConsentDeniedException; +import org.springframework.ai.tool.definition.ToolDefinition; +import org.springframework.ai.tool.metadata.ToolMetadata; +import org.springframework.util.Assert; + +/** + * A decorator for {@link ToolCallback} that enforces consent requirements before + * delegating to the actual tool implementation. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +public class ConsentAwareToolCallback implements ToolCallback { + + private static final Pattern PLACEHOLDER_PATTERN = Pattern.compile("\\{([^}]+)\\}"); + + private final ToolCallback delegate; + + private final ConsentManager consentManager; + + private final RequiresConsent requiresConsent; + + /** + * Creates a new consent-aware tool callback. + * @param delegate the actual tool callback to delegate to + * @param consentManager the consent manager for handling consent requests + * @param requiresConsent the consent requirements annotation + */ + public ConsentAwareToolCallback(ToolCallback delegate, ConsentManager consentManager, + RequiresConsent requiresConsent) { + Assert.notNull(delegate, "delegate must not be null"); + Assert.notNull(consentManager, "consentManager must not be null"); + Assert.notNull(requiresConsent, "requiresConsent must not be null"); + this.delegate = delegate; + this.consentManager = consentManager; + this.requiresConsent = requiresConsent; + } + + private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper(); + + @Override + public String call(String toolInput) { + // Parse input JSON to get parameters + Map parameters; + try { + parameters = OBJECT_MAPPER.readValue(toolInput, new TypeReference>() { + }); + } + catch (Exception e) { + parameters = Map.of(); + } + String toolName = this.delegate.getToolDefinition().name(); + + // Check if consent was already granted based on consent level + if (this.consentManager.hasValidConsent(toolName, this.requiresConsent.level(), + this.requiresConsent.categories())) { + return this.delegate.call(toolInput); + } + + // Prepare consent message with parameter substitution + String message = prepareConsentMessage(this.requiresConsent.message(), parameters); + + // Request consent + boolean consentGranted = this.consentManager.requestConsent(toolName, message, this.requiresConsent.level(), + this.requiresConsent.categories(), parameters); + + if (!consentGranted) { + throw new ConsentDeniedException(String.format("User denied consent for tool '%s' execution", toolName)); + } + + // Execute the tool if consent was granted + return this.delegate.call(toolInput); + } + + @Override + public ToolDefinition getToolDefinition() { + return this.delegate.getToolDefinition(); + } + + @Override + public ToolMetadata getToolMetadata() { + return this.delegate.getToolMetadata(); + } + + /** + * Prepares the consent message by replacing placeholders with actual parameter + * values. + * @param template the message template with placeholders + * @param parameters the parameters to substitute + * @return the prepared message + */ + private String prepareConsentMessage(String template, Map parameters) { + if (parameters == null || parameters.isEmpty()) { + return template; + } + + Matcher matcher = PLACEHOLDER_PATTERN.matcher(template); + StringBuffer result = new StringBuffer(); + + while (matcher.find()) { + String paramName = matcher.group(1); + Object value = parameters.get(paramName); + String replacement = value != null ? String.valueOf(value) : "{" + paramName + "}"; + matcher.appendReplacement(result, Matcher.quoteReplacement(replacement)); + } + matcher.appendTail(result); + + return result.toString(); + } + + /** + * Returns the underlying delegate tool callback. + * @return the delegate tool callback + */ + public ToolCallback getDelegate() { + return this.delegate; + } + + /** + * Returns the consent requirements for this tool. + * @return the consent requirements + */ + public RequiresConsent getRequiresConsent() { + return this.requiresConsent; + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentChecker.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentChecker.java new file mode 100644 index 00000000000..c824215b82e --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentChecker.java @@ -0,0 +1,36 @@ +/* + * 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.consent; + +/** + * Interface for checking user consent before tool execution. Implementations can provide + * different strategies for obtaining and validating user consent. + * + * @author Assistant + * @since 1.0.0 + */ +@FunctionalInterface +public interface ConsentChecker { + + /** + * Check if consent is granted for the given context. + * @param context the consent context containing tool information + * @return true if consent is granted, false otherwise + */ + boolean checkConsent(ConsentContext context); + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentContext.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentContext.java new file mode 100644 index 00000000000..17ca19b8df6 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentContext.java @@ -0,0 +1,139 @@ +/* + * 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.consent; + +import org.springframework.ai.chat.model.ToolContext; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.lang.Nullable; + +/** + * Context information for consent checking. Contains all relevant information needed to + * make consent decisions. + * + * @author Assistant + * @since 1.0.0 + */ +public class ConsentContext { + + private final String toolName; + + private final String toolInput; + + @Nullable + private final ToolContext toolContext; + + @Nullable + private final RequiresConsent requiresConsent; + + @Nullable + private final String userId; + + @Nullable + private final String sessionId; + + private ConsentContext(Builder builder) { + this.toolName = builder.toolName; + this.toolInput = builder.toolInput; + this.toolContext = builder.toolContext; + this.requiresConsent = builder.requiresConsent; + this.userId = builder.userId; + this.sessionId = builder.sessionId; + } + + public String getToolName() { + return toolName; + } + + public String getToolInput() { + return toolInput; + } + + @Nullable + public ToolContext getToolContext() { + return toolContext; + } + + @Nullable + public RequiresConsent getRequiresConsent() { + return requiresConsent; + } + + @Nullable + public String getUserId() { + return userId; + } + + @Nullable + public String getSessionId() { + return sessionId; + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + + private String toolName; + + private String toolInput; + + private ToolContext toolContext; + + private RequiresConsent requiresConsent; + + private String userId; + + private String sessionId; + + public Builder toolName(String toolName) { + this.toolName = toolName; + return this; + } + + public Builder toolInput(String toolInput) { + this.toolInput = toolInput; + return this; + } + + public Builder toolContext(ToolContext toolContext) { + this.toolContext = toolContext; + return this; + } + + public Builder requiresConsent(RequiresConsent requiresConsent) { + this.requiresConsent = requiresConsent; + return this; + } + + public Builder userId(String userId) { + this.userId = userId; + return this; + } + + public Builder sessionId(String sessionId) { + this.sessionId = sessionId; + return this; + } + + public ConsentContext build() { + return new ConsentContext(this); + } + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentManager.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentManager.java new file mode 100644 index 00000000000..fbb25eb40ae --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentManager.java @@ -0,0 +1,67 @@ +/* + * 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.consent; + +import java.util.Map; + +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; + +/** + * Strategy interface for managing user consent for tool execution. Implementations can + * provide different consent mechanisms such as UI-based prompts, API calls, or + * configuration-based consent. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +public interface ConsentManager { + + /** + * Requests consent for executing a tool with the given parameters. + * @param toolName the name of the tool requiring consent + * @param message the consent message (may contain placeholders) + * @param level the consent level required + * @param categories the categories associated with this consent request + * @param parameters the actual parameters that will be passed to the tool + * @return {@code true} if consent is granted, {@code false} otherwise + */ + boolean requestConsent(String toolName, String message, ConsentLevel level, String[] categories, + Map parameters); + + /** + * Checks if consent has already been granted for a tool based on the consent level. + * This method is called before {@link #requestConsent} to avoid unnecessary prompts. + * @param toolName the name of the tool + * @param level the consent level + * @param categories the categories associated with this tool + * @return {@code true} if consent was previously granted and is still valid + */ + boolean hasValidConsent(String toolName, ConsentLevel level, String[] categories); + + /** + * Revokes any stored consent for the specified tool. + * @param toolName the name of the tool + * @param categories the categories to revoke consent for (empty to revoke all) + */ + void revokeConsent(String toolName, String[] categories); + + /** + * Clears all stored consents. Typically called at session end or on user request. + */ + void clearAllConsents(); + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentRequiredException.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentRequiredException.java new file mode 100644 index 00000000000..2508ad995d6 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/ConsentRequiredException.java @@ -0,0 +1,35 @@ +/* + * 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.consent; + +/** + * Exception thrown when a tool requires user consent but consent was not granted. + * + * @author Assistant + * @since 1.0.0 + */ +public class ConsentRequiredException extends RuntimeException { + + public ConsentRequiredException(String message) { + super(message); + } + + public ConsentRequiredException(String message, Throwable cause) { + super(message, cause); + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentChecker.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentChecker.java new file mode 100644 index 00000000000..fdc479837f3 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentChecker.java @@ -0,0 +1,112 @@ +/* + * 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.consent; + +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Default implementation of {@link ConsentChecker} that maintains an in-memory store of + * granted consents. + * + * @author Assistant + * @since 1.0.0 + */ +public class DefaultConsentChecker implements ConsentChecker { + + private static final Logger logger = LoggerFactory.getLogger(DefaultConsentChecker.class); + + private final Set grantedConsents = ConcurrentHashMap.newKeySet(); + + private final ConsentRequestHandler consentRequestHandler; + + public DefaultConsentChecker(ConsentRequestHandler consentRequestHandler) { + this.consentRequestHandler = consentRequestHandler; + } + + @Override + public boolean checkConsent(ConsentContext context) { + String consentKey = buildConsentKey(context); + + // Check if consent already granted + if (grantedConsents.contains(consentKey)) { + logger.debug("Consent already granted for tool: {}", context.getToolName()); + return true; + } + + // Request consent + boolean consentGranted = consentRequestHandler.requestConsent(context); + + if (consentGranted) { + // Store consent if granted + if (context.getRequiresConsent() != null + && context.getRequiresConsent().level() != RequiresConsent.ConsentLevel.ONE_TIME) { + grantedConsents.add(consentKey); + } + logger.info("Consent granted for tool: {}", context.getToolName()); + } + else { + logger.warn("Consent denied for tool: {}", context.getToolName()); + } + + return consentGranted; + } + + private String buildConsentKey(ConsentContext context) { + StringBuilder key = new StringBuilder(); + key.append(context.getToolName()); + + if (context.getUserId() != null) { + key.append(":").append(context.getUserId()); + } + + if (context.getSessionId() != null && context.getRequiresConsent() != null + && context.getRequiresConsent().level() == RequiresConsent.ConsentLevel.SESSION) { + key.append(":").append(context.getSessionId()); + } + + return key.toString(); + } + + public void revokeConsent(String toolName, String userId) { + String prefix = toolName + ":" + userId; + grantedConsents.removeIf(key -> key.startsWith(prefix)); + } + + public void clearAllConsents() { + grantedConsents.clear(); + } + + /** + * Interface for handling consent requests from users. + */ + @FunctionalInterface + public interface ConsentRequestHandler { + + /** + * Request consent from the user. + * @param context the consent context + * @return true if consent is granted, false otherwise + */ + boolean requestConsent(ConsentContext context); + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentManager.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentManager.java new file mode 100644 index 00000000000..b75f70a624d --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/DefaultConsentManager.java @@ -0,0 +1,156 @@ +/* + * 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.consent; + +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.BiFunction; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; + +/** + * Default implementation of {@link ConsentManager} that uses a configurable consent + * handler function and stores consent decisions in memory. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +public class DefaultConsentManager implements ConsentManager { + + private static final Logger logger = LoggerFactory.getLogger(DefaultConsentManager.class); + + private final BiFunction, Boolean> consentHandler; + + private final Map consentStore = new ConcurrentHashMap<>(); + + /** + * Creates a default consent manager with the given consent handler. + * @param consentHandler function that takes a consent message and parameters, returns + * true if consent is granted + */ + public DefaultConsentManager(BiFunction, Boolean> consentHandler) { + this.consentHandler = consentHandler; + } + + /** + * Creates a default consent manager that always denies consent. Useful for testing or + * when consent should be managed externally. + */ + public DefaultConsentManager() { + this((message, params) -> { + logger.warn("No consent handler configured. Denying consent for: {}", message); + return false; + }); + } + + @Override + public boolean requestConsent(String toolName, String message, ConsentLevel level, String[] categories, + Map parameters) { + + logger.debug("Requesting consent for tool '{}' with message: {}", toolName, message); + + // Check if we already have valid consent + if (hasValidConsent(toolName, level, categories)) { + logger.debug("Valid consent already exists for tool '{}'", toolName); + return true; + } + + // Request new consent + boolean granted = this.consentHandler.apply(message, parameters); + + if (granted) { + // Store consent decision based on level + if (level != ConsentLevel.EVERY_TIME) { + String key = createConsentKey(toolName, categories); + this.consentStore.put(key, new ConsentRecord(level, System.currentTimeMillis())); + logger.debug("Stored consent for tool '{}' with level '{}'", toolName, level); + } + } + + return granted; + } + + @Override + public boolean hasValidConsent(String toolName, ConsentLevel level, String[] categories) { + if (level == ConsentLevel.EVERY_TIME) { + return false; // Always require new consent + } + + String key = createConsentKey(toolName, categories); + ConsentRecord record = this.consentStore.get(key); + + if (record == null) { + return false; + } + + // For REMEMBER level, consent is valid indefinitely + // For SESSION level, we could implement session timeout logic here + return true; + } + + @Override + public void revokeConsent(String toolName, String[] categories) { + if (categories == null || categories.length == 0) { + // Revoke all consents for this tool + this.consentStore.entrySet().removeIf(entry -> entry.getKey().startsWith(toolName + ":")); + } + else { + // Revoke specific category consent + String key = createConsentKey(toolName, categories); + this.consentStore.remove(key); + } + logger.debug("Revoked consent for tool '{}'", toolName); + } + + @Override + public void clearAllConsents() { + this.consentStore.clear(); + logger.debug("Cleared all stored consents"); + } + + /** + * Creates a unique key for storing consent based on tool name and categories. + */ + private String createConsentKey(String toolName, String[] categories) { + if (categories == null || categories.length == 0) { + return toolName + ":default"; + } + String sortedCategories = Arrays.stream(categories).sorted().reduce((a, b) -> a + "," + b).orElse(""); + return toolName + ":" + sortedCategories; + } + + /** + * Internal record for storing consent information. + */ + private static class ConsentRecord { + + final ConsentLevel level; + + final long timestamp; + + ConsentRecord(ConsentLevel level, long timestamp) { + this.level = level; + this.timestamp = timestamp; + } + + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/exception/ConsentDeniedException.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/exception/ConsentDeniedException.java new file mode 100644 index 00000000000..1494658413e --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/exception/ConsentDeniedException.java @@ -0,0 +1,45 @@ +/* + * 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.consent.exception; + +/** + * Exception thrown when user consent is required but denied for tool execution. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +public class ConsentDeniedException extends RuntimeException { + + /** + * Constructs a new consent denied exception with the specified detail message. + * @param message the detail message + */ + public ConsentDeniedException(String message) { + super(message); + } + + /** + * Constructs a new consent denied exception with the specified detail message and + * cause. + * @param message the detail message + * @param cause the cause + */ + public ConsentDeniedException(String message, Throwable cause) { + super(message, cause); + } + +} diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/package-info.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/package-info.java new file mode 100644 index 00000000000..68cfe2d8ff7 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/consent/package-info.java @@ -0,0 +1,49 @@ +/* + * 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. + */ + +/** + * Consent management framework for Spring AI tool execution. + * + *

+ * This package provides a comprehensive consent management system that allows tools to + * require user approval before execution. Key components include: + * + *

    + *
  • {@link org.springframework.ai.tool.annotation.RequiresConsent} - Annotation to mark + * tool methods that require consent
    • + *
  • {@link org.springframework.ai.tool.consent.ConsentManager} - Strategy interface for + * managing consent requests and decisions
    • + *
  • {@link org.springframework.ai.tool.consent.ConsentAwareToolCallback} - Decorator + * that enforces consent requirements before tool execution
    • + *
  • {@link org.springframework.ai.tool.consent.DefaultConsentManager} - Default + * implementation with configurable consent handlers
    • + *
+ * + *

+ * Example usage: 

{@code
+ * @Tool(description = "Deletes a record")
+ * @RequiresConsent(
+ *     message = "Delete record {id}? This cannot be undone.",
+ *     level = ConsentLevel.EVERY_TIME
+ * )
+ * public void deleteRecord(String id) {
+ *     // Implementation
+ * }
+ * }
+ * + * @since 1.0.0 + */ +package org.springframework.ai.tool.consent; diff --git a/spring-ai-model/src/main/java/org/springframework/ai/tool/method/ConsentAwareMethodToolCallbackProvider.java b/spring-ai-model/src/main/java/org/springframework/ai/tool/method/ConsentAwareMethodToolCallbackProvider.java new file mode 100644 index 00000000000..969f46157c0 --- /dev/null +++ b/spring-ai-model/src/main/java/org/springframework/ai/tool/method/ConsentAwareMethodToolCallbackProvider.java @@ -0,0 +1,178 @@ +/* + * 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.method; + +import java.lang.reflect.Method; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.function.Consumer; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.ToolCallbackProvider; +import org.springframework.ai.tool.annotation.Tool; +import org.springframework.ai.tool.consent.ConsentAwareToolCallback; +import org.springframework.ai.tool.consent.ConsentChecker; +import org.springframework.ai.tool.metadata.ToolMetadata; +import org.springframework.ai.tool.support.ToolDefinitions; +import org.springframework.ai.tool.support.ToolUtils; +import org.springframework.aop.support.AopUtils; +import org.springframework.core.annotation.AnnotationUtils; +import org.springframework.util.Assert; +import org.springframework.util.ClassUtils; +import org.springframework.util.ReflectionUtils; + +/** + * A {@link ToolCallbackProvider} that builds {@link ToolCallback} instances from + * {@link Tool}-annotated methods with consent awareness support. + * + * @author Assistant + * @since 1.0.0 + */ +public final class ConsentAwareMethodToolCallbackProvider implements ToolCallbackProvider { + + private static final Logger logger = LoggerFactory.getLogger(ConsentAwareMethodToolCallbackProvider.class); + + private final List toolObjects; + + private final ConsentChecker consentChecker; + + private ConsentAwareMethodToolCallbackProvider(List toolObjects, ConsentChecker consentChecker) { + Assert.notNull(toolObjects, "toolObjects cannot be null"); + Assert.notNull(consentChecker, "consentChecker cannot be null"); + Assert.noNullElements(toolObjects, "toolObjects cannot contain null elements"); + assertToolAnnotatedMethodsPresent(toolObjects); + this.toolObjects = toolObjects; + this.consentChecker = consentChecker; + validateToolCallbacks(getToolCallbacks()); + } + + private void assertToolAnnotatedMethodsPresent(List toolObjects) { + for (Object toolObject : toolObjects) { + List toolMethods = Stream + .of(ReflectionUtils.getDeclaredMethods( + AopUtils.isAopProxy(toolObject) ? AopUtils.getTargetClass(toolObject) : toolObject.getClass())) + .filter(this::isToolAnnotatedMethod) + .filter(toolMethod -> !isFunctionalType(toolMethod)) + .toList(); + + if (toolMethods.isEmpty()) { + throw new IllegalStateException("No @Tool annotated methods found in " + toolObject + "." + + "Did you mean to pass a ToolCallback or ToolCallbackProvider? If so, you have to use .toolCallbacks() instead of .tool()"); + } + } + } + + @Override + public ToolCallback[] getToolCallbacks() { + var toolCallbacks = this.toolObjects.stream() + .map(toolObject -> Stream + .of(ReflectionUtils.getDeclaredMethods( + AopUtils.isAopProxy(toolObject) ? AopUtils.getTargetClass(toolObject) : toolObject.getClass())) + .filter(this::isToolAnnotatedMethod) + .filter(toolMethod -> !isFunctionalType(toolMethod)) + .map(toolMethod -> { + ToolCallback baseCallback = MethodToolCallback.builder() + .toolDefinition(ToolDefinitions.from(toolMethod)) + .toolMetadata(ToolMetadata.from(toolMethod)) + .toolMethod(toolMethod) + .toolObject(toolObject) + .toolCallResultConverter(ToolUtils.getToolCallResultConverter(toolMethod)) + .build(); + + // Wrap with consent checking if needed + if (ToolUtils.requiresConsent(toolMethod)) { + return new ConsentAwareToolCallback(baseCallback, this.consentChecker, true); + } + return baseCallback; + }) + .toArray(ToolCallback[]::new)) + .flatMap(Stream::of) + .toArray(ToolCallback[]::new); + + validateToolCallbacks(toolCallbacks); + + return toolCallbacks; + } + + private boolean isFunctionalType(Method toolMethod) { + var isFunction = ClassUtils.isAssignable(Function.class, toolMethod.getReturnType()) + || ClassUtils.isAssignable(Supplier.class, toolMethod.getReturnType()) + || ClassUtils.isAssignable(Consumer.class, toolMethod.getReturnType()); + + if (isFunction) { + logger.warn("Method {} is annotated with @Tool but returns a functional type. " + + "This is not supported and the method will be ignored.", toolMethod.getName()); + } + + return isFunction; + } + + private boolean isToolAnnotatedMethod(Method method) { + Tool annotation = AnnotationUtils.findAnnotation(method, Tool.class); + return Objects.nonNull(annotation); + } + + private void validateToolCallbacks(ToolCallback[] toolCallbacks) { + List duplicateToolNames = ToolUtils.getDuplicateToolNames(toolCallbacks); + if (!duplicateToolNames.isEmpty()) { + throw new IllegalStateException("Multiple tools with the same name (%s) found in sources: %s".formatted( + String.join(", ", duplicateToolNames), + this.toolObjects.stream().map(o -> o.getClass().getName()).collect(Collectors.joining(", ")))); + } + } + + public static Builder builder() { + return new Builder(); + } + + public static final class Builder { + + private List toolObjects; + + private ConsentChecker consentChecker; + + private Builder() { + } + + public Builder toolObjects(Object... toolObjects) { + Assert.notNull(toolObjects, "toolObjects cannot be null"); + this.toolObjects = Arrays.asList(toolObjects); + return this; + } + + public Builder consentChecker(ConsentChecker consentChecker) { + Assert.notNull(consentChecker, "consentChecker cannot be null"); + this.consentChecker = consentChecker; + return this; + } + + public ConsentAwareMethodToolCallbackProvider build() { + Assert.notNull(this.consentChecker, "consentChecker must be set"); + return new ConsentAwareMethodToolCallbackProvider(this.toolObjects, this.consentChecker); + } + + } + +} diff --git a/spring-ai-model/src/test/java/org/springframework/ai/model/tool/DefaultToolCallingManagerConsentTests.java b/spring-ai-model/src/test/java/org/springframework/ai/model/tool/DefaultToolCallingManagerConsentTests.java new file mode 100644 index 00000000000..06aadf9adbc --- /dev/null +++ b/spring-ai-model/src/test/java/org/springframework/ai/model/tool/DefaultToolCallingManagerConsentTests.java @@ -0,0 +1,266 @@ +/* + * 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.model.tool; + +import java.util.List; +import java.util.Map; + +import io.micrometer.observation.ObservationRegistry; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; +import org.springframework.ai.tool.consent.ConsentAwareToolCallback; +import org.springframework.ai.tool.consent.ConsentManager; +import org.springframework.ai.tool.consent.exception.ConsentDeniedException; +import org.springframework.ai.tool.definition.DefaultToolDefinition; +import org.springframework.ai.tool.definition.ToolDefinition; +import org.springframework.ai.tool.execution.DefaultToolExecutionExceptionProcessor; +import org.springframework.ai.tool.metadata.DefaultToolMetadata; +import org.springframework.ai.tool.metadata.ToolMetadata; +import org.springframework.ai.tool.resolution.ToolCallbackResolver; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link DefaultToolCallingManager} with consent management. + * + * @author Hyunjoon Park + * @since 1.0.0 + */ +class DefaultToolCallingManagerConsentTests { + + private DefaultToolCallingManager toolCallingManager; + + private ConsentManager consentManager; + + private ToolCallback mockToolCallback; + + private ConsentAwareToolCallback consentAwareToolCallback; + + @BeforeEach + void setUp() { + ObservationRegistry observationRegistry = ObservationRegistry.create(); + ToolCallbackResolver toolCallbackResolver = Mockito.mock(ToolCallbackResolver.class); + DefaultToolExecutionExceptionProcessor exceptionProcessor = DefaultToolExecutionExceptionProcessor.builder() + .build(); + + this.toolCallingManager = new DefaultToolCallingManager(observationRegistry, toolCallbackResolver, + exceptionProcessor); + + // Set up mock tool callback + this.mockToolCallback = Mockito.mock(ToolCallback.class); + ToolDefinition toolDefinition = DefaultToolDefinition.builder() + .name("deleteBook") + .description("Delete a book") + .inputSchema("{\"type\":\"object\",\"properties\":{\"bookId\":{\"type\":\"string\"}}}") + .build(); + ToolMetadata toolMetadata = DefaultToolMetadata.builder().build(); + when(this.mockToolCallback.getToolDefinition()).thenReturn(toolDefinition); + when(this.mockToolCallback.getToolMetadata()).thenReturn(toolMetadata); + + // Set up consent manager + this.consentManager = Mockito.mock(ConsentManager.class); + + // Create mock RequiresConsent annotation + RequiresConsent requiresConsent = Mockito.mock(RequiresConsent.class); + when(requiresConsent.message()).thenReturn("Delete book {bookId}?"); + when(requiresConsent.level()).thenReturn(ConsentLevel.EVERY_TIME); + when(requiresConsent.categories()).thenReturn(new String[0]); + + // Create consent-aware wrapper + this.consentAwareToolCallback = new ConsentAwareToolCallback(this.mockToolCallback, this.consentManager, + requiresConsent); + } + + @Test + void testManualExecutionWithConsentGranted() { + // Given + // ConsentAwareToolCallback will first check hasValidConsent, then call + // requestConsent if needed + // For this test, we'll make hasValidConsent return false to trigger + // requestConsent + when(this.consentManager.hasValidConsent(anyString(), any(ConsentLevel.class), any(String[].class))) + .thenReturn(false); + when(this.consentManager.requestConsent(anyString(), anyString(), any(ConsentLevel.class), any(String[].class), + any(Map.class))) + .thenReturn(true); + when(this.mockToolCallback.call(anyString(), any())).thenReturn("Book deleted"); + + List toolCallbacks = List.of(this.consentAwareToolCallback); + ToolCallingChatOptions chatOptions = ToolCallingChatOptions.builder().toolCallbacks(toolCallbacks).build(); + + UserMessage userMessage = new UserMessage("Delete book with ID 123"); + Prompt prompt = new Prompt(List.of(userMessage), chatOptions); + + AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall("1", "tool-call", "deleteBook", + "{\"bookId\":\"123\"}"); + AssistantMessage assistantMessage = new AssistantMessage("I'll delete the book.", Map.of(), List.of(toolCall)); + + Generation generation = new Generation(assistantMessage); + ChatResponse chatResponse = new ChatResponse(List.of(generation)); + + // When + ToolExecutionResult result = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + + // Then + assertThat(result).isNotNull(); + assertThat(result.conversationHistory()).hasSize(3); // user, assistant, tool + // response + // Verify consent was requested + verify(this.consentManager, times(1)).hasValidConsent(anyString(), any(ConsentLevel.class), + any(String[].class)); + verify(this.consentManager, times(1)).requestConsent(anyString(), anyString(), any(ConsentLevel.class), + any(String[].class), any(Map.class)); + verify(this.mockToolCallback, times(1)).call(anyString(), any()); + } + + @Test + void testManualExecutionWithConsentDenied() { + // Given + // ConsentAwareToolCallback will first check hasValidConsent, then call + // requestConsent if needed + when(this.consentManager.hasValidConsent(anyString(), any(ConsentLevel.class), any(String[].class))) + .thenReturn(false); + when(this.consentManager.requestConsent(anyString(), anyString(), any(ConsentLevel.class), any(String[].class), + any(Map.class))) + .thenReturn(false); + + List toolCallbacks = List.of(this.consentAwareToolCallback); + ToolCallingChatOptions chatOptions = ToolCallingChatOptions.builder().toolCallbacks(toolCallbacks).build(); + + UserMessage userMessage = new UserMessage("Delete book with ID 123"); + Prompt prompt = new Prompt(List.of(userMessage), chatOptions); + + AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall("1", "tool-call", "deleteBook", + "{\"bookId\":\"123\"}"); + AssistantMessage assistantMessage = new AssistantMessage("I'll delete the book.", Map.of(), List.of(toolCall)); + + Generation generation = new Generation(assistantMessage); + ChatResponse chatResponse = new ChatResponse(List.of(generation)); + + // When & Then + assertThatThrownBy(() -> this.toolCallingManager.executeToolCalls(prompt, chatResponse)) + .isInstanceOf(ConsentDeniedException.class) + .hasMessageContaining("User denied consent for tool"); + + // Verify consent was requested but denied + verify(this.consentManager, times(1)).hasValidConsent(anyString(), any(ConsentLevel.class), + any(String[].class)); + verify(this.consentManager, times(1)).requestConsent(anyString(), anyString(), any(ConsentLevel.class), + any(String[].class), any(Map.class)); + verify(this.mockToolCallback, times(0)).call(anyString(), any()); + } + + @Test + void testManualExecutionWithNonConsentAwareToolCallback() { + // Given + when(this.mockToolCallback.call(anyString(), any())).thenReturn("Book deleted"); + + List toolCallbacks = List.of(this.mockToolCallback); // Regular + // callback, + // not + // consent-aware + ToolCallingChatOptions chatOptions = ToolCallingChatOptions.builder().toolCallbacks(toolCallbacks).build(); + + UserMessage userMessage = new UserMessage("Delete book with ID 123"); + Prompt prompt = new Prompt(List.of(userMessage), chatOptions); + + AssistantMessage.ToolCall toolCall = new AssistantMessage.ToolCall("1", "tool-call", "deleteBook", + "{\"bookId\":\"123\"}"); + AssistantMessage assistantMessage = new AssistantMessage("I'll delete the book.", Map.of(), List.of(toolCall)); + + Generation generation = new Generation(assistantMessage); + ChatResponse chatResponse = new ChatResponse(List.of(generation)); + + // When + ToolExecutionResult result = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + + // Then + assertThat(result).isNotNull(); + assertThat(result.conversationHistory()).hasSize(3); + verify(this.mockToolCallback, times(1)).call(anyString(), any()); + // ConsentManager should not be called for non-consent-aware callbacks + } + + @Test + void testManualExecutionWithMixedToolCallbacks() { + // Given + ToolCallback regularCallback = Mockito.mock(ToolCallback.class); + ToolDefinition regularToolDef = DefaultToolDefinition.builder() + .name("getBook") + .description("Get a book") + .inputSchema("{\"type\":\"object\",\"properties\":{\"bookId\":{\"type\":\"string\"}}}") + .build(); + when(regularCallback.getToolDefinition()).thenReturn(regularToolDef); + when(regularCallback.getToolMetadata()).thenReturn(DefaultToolMetadata.builder().build()); + when(regularCallback.call(anyString(), any())).thenReturn("Book found"); + + // For the consent-aware callback + when(this.consentManager.hasValidConsent(anyString(), any(ConsentLevel.class), any(String[].class))) + .thenReturn(false); + when(this.consentManager.requestConsent(anyString(), anyString(), any(ConsentLevel.class), any(String[].class), + any(Map.class))) + .thenReturn(true); + when(this.mockToolCallback.call(anyString(), any())).thenReturn("Book deleted"); + + List toolCallbacks = List.of(regularCallback, this.consentAwareToolCallback); + ToolCallingChatOptions chatOptions = ToolCallingChatOptions.builder().toolCallbacks(toolCallbacks).build(); + + UserMessage userMessage = new UserMessage("Get and delete book with ID 123"); + Prompt prompt = new Prompt(List.of(userMessage), chatOptions); + + AssistantMessage.ToolCall getCall = new AssistantMessage.ToolCall("1", "tool-call", "getBook", + "{\"bookId\":\"123\"}"); + AssistantMessage.ToolCall deleteCall = new AssistantMessage.ToolCall("2", "tool-call", "deleteBook", + "{\"bookId\":\"123\"}"); + AssistantMessage assistantMessage = new AssistantMessage("I'll get and delete the book.", Map.of(), + List.of(getCall, deleteCall)); + + Generation generation = new Generation(assistantMessage); + ChatResponse chatResponse = new ChatResponse(List.of(generation)); + + // When + ToolExecutionResult result = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + + // Then + assertThat(result).isNotNull(); + assertThat(result.conversationHistory()).hasSize(3); + verify(regularCallback, times(1)).call(anyString(), any()); + // Verify consent was requested for the consent-aware callback only + verify(this.consentManager, times(1)).hasValidConsent(anyString(), any(ConsentLevel.class), + any(String[].class)); + verify(this.consentManager, times(1)).requestConsent(anyString(), anyString(), any(ConsentLevel.class), + any(String[].class), any(Map.class)); + verify(this.mockToolCallback, times(1)).call(anyString(), any()); + } + +} \ No newline at end of file diff --git a/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTest.java b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTest.java new file mode 100644 index 00000000000..d5227d57deb --- /dev/null +++ b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTest.java @@ -0,0 +1,152 @@ +/* + * 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.consent; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.annotation.Tool; +import org.springframework.ai.tool.method.ConsentAwareMethodToolCallbackProvider; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for consent-aware tool execution. + * + * @author Assistant + */ +class ConsentAwareToolCallbackTest { + + @Test + void testToolWithConsentGranted() { + // Create a consent checker that always grants consent + ConsentChecker alwaysGrantChecker = context -> true; + + ConsentAwareMethodToolCallbackProvider provider = ConsentAwareMethodToolCallbackProvider.builder() + .toolObjects(new ConsentRequiredTools()) + .consentChecker(alwaysGrantChecker) + .build(); + + ToolCallback[] callbacks = provider.getToolCallbacks(); + assertThat(callbacks).hasSize(2); + + // Find the sendEmail tool + ToolCallback sendEmailTool = findToolByName(callbacks, "sendEmail"); + assertThat(sendEmailTool).isNotNull(); + + // Execute should succeed when consent is granted + String result = sendEmailTool.call("{\"to\":\"test@example.com\",\"subject\":\"Test\",\"body\":\"Hello\"}"); + assertThat(result).isEqualTo("Email sent to test@example.com"); + } + + @Test + void testToolWithConsentDenied() { + // Create a consent checker that always denies consent + ConsentChecker alwaysDenyChecker = context -> false; + + ConsentAwareMethodToolCallbackProvider provider = ConsentAwareMethodToolCallbackProvider.builder() + .toolObjects(new ConsentRequiredTools()) + .consentChecker(alwaysDenyChecker) + .build(); + + ToolCallback[] callbacks = provider.getToolCallbacks(); + ToolCallback sendEmailTool = findToolByName(callbacks, "sendEmail"); + + // Execute should fail when consent is denied + assertThatThrownBy( + () -> sendEmailTool.call("{\"to\":\"test@example.com\",\"subject\":\"Test\",\"body\":\"Hello\"}")) + .isInstanceOf(ConsentRequiredException.class) + .hasMessageContaining("User consent required for tool: sendEmail"); + } + + @Test + void testToolWithoutConsentAnnotation() { + // Tools without @RequiresConsent should work normally + ConsentChecker neverCalledChecker = context -> { + throw new AssertionError("Consent checker should not be called for tools without @RequiresConsent"); + }; + + ConsentAwareMethodToolCallbackProvider provider = ConsentAwareMethodToolCallbackProvider.builder() + .toolObjects(new ConsentRequiredTools()) + .consentChecker(neverCalledChecker) + .build(); + + ToolCallback[] callbacks = provider.getToolCallbacks(); + ToolCallback getTimeTool = findToolByName(callbacks, "getCurrentTime"); + + // Should execute without consent check + String result = getTimeTool.call("{}"); + assertThat(result).contains("Current time:"); + } + + @Test + void testDefaultConsentChecker() { + // Test with the default consent checker + DefaultConsentChecker.ConsentRequestHandler handler = context -> { + // Simulate user interaction - grant consent for sendEmail only + return "sendEmail".equals(context.getToolName()); + }; + + DefaultConsentChecker consentChecker = new DefaultConsentChecker(handler); + + ConsentAwareMethodToolCallbackProvider provider = ConsentAwareMethodToolCallbackProvider.builder() + .toolObjects(new ConsentRequiredTools()) + .consentChecker(consentChecker) + .build(); + + ToolCallback[] callbacks = provider.getToolCallbacks(); + ToolCallback sendEmailTool = findToolByName(callbacks, "sendEmail"); + + // First call should request and grant consent + String result = sendEmailTool.call("{\"to\":\"test@example.com\",\"subject\":\"Test\",\"body\":\"Hello\"}"); + assertThat(result).isEqualTo("Email sent to test@example.com"); + + // Second call should use cached consent (for SESSION level) + result = sendEmailTool.call("{\"to\":\"another@example.com\",\"subject\":\"Test2\",\"body\":\"Hello again\"}"); + assertThat(result).isEqualTo("Email sent to another@example.com"); + } + + private ToolCallback findToolByName(ToolCallback[] callbacks, String name) { + for (ToolCallback callback : callbacks) { + if (callback.getToolDefinition().name().equals(name)) { + return callback; + } + } + return null; + } + + /** + * Test class with tools that require consent. + */ + static class ConsentRequiredTools { + + @Tool(description = "Send an email to a recipient") + @RequiresConsent(reason = "This tool will send an email on your behalf", + level = RequiresConsent.ConsentLevel.SESSION, categories = { "communication", "email" }) + public String sendEmail(String to, String subject, String body) { + return "Email sent to " + to; + } + + @Tool(description = "Get the current time") + public String getCurrentTime() { + return "Current time: " + System.currentTimeMillis(); + } + + } + +} diff --git a/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTests.java b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTests.java new file mode 100644 index 00000000000..57cf0fa17ba --- /dev/null +++ b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/ConsentAwareToolCallbackTests.java @@ -0,0 +1,172 @@ +/* + * 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.consent; + +import java.util.HashMap; +import java.util.Map; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; +import org.springframework.ai.tool.consent.exception.ConsentDeniedException; +import org.springframework.ai.tool.definition.ToolDefinition; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link ConsentAwareToolCallback}. + * + * @author Hyunjoon Park + */ +class ConsentAwareToolCallbackTests { + + @Mock + private ToolCallback delegate; + + @Mock + private ConsentManager consentManager; + + @Mock + private RequiresConsent requiresConsent; + + private ConsentAwareToolCallback consentAwareCallback; + + @BeforeEach + void setUp() { + MockitoAnnotations.openMocks(this); + when(requiresConsent.message()).thenReturn("Do you approve this action?"); + when(requiresConsent.level()).thenReturn(ConsentLevel.EVERY_TIME); + when(requiresConsent.categories()).thenReturn(new String[0]); + + consentAwareCallback = new ConsentAwareToolCallback(delegate, consentManager, requiresConsent); + } + + @Test + void callWithConsentGranted() { + // Given + Map parameters = Map.of("param1", "value1"); + Object expectedResult = "Tool executed successfully"; + + when(delegate.getName()).thenReturn("testTool"); + when(consentManager.hasValidConsent("testTool", ConsentLevel.EVERY_TIME, new String[0])).thenReturn(false); + when(consentManager.requestConsent(eq("testTool"), any(), any(), any(), eq(parameters))).thenReturn(true); + when(delegate.call(parameters)).thenReturn(expectedResult); + + // When + Object result = consentAwareCallback.call(parameters); + + // Then + assertThat(result).isEqualTo(expectedResult); + verify(delegate).call(parameters); + verify(consentManager).requestConsent(eq("testTool"), any(), any(), any(), eq(parameters)); + } + + @Test + void callWithConsentDenied() { + // Given + Map parameters = Map.of("param1", "value1"); + + when(delegate.getName()).thenReturn("testTool"); + when(consentManager.hasValidConsent("testTool", ConsentLevel.EVERY_TIME, new String[0])).thenReturn(false); + when(consentManager.requestConsent(eq("testTool"), any(), any(), any(), eq(parameters))).thenReturn(false); + + // When/Then + assertThatThrownBy(() -> consentAwareCallback.call(parameters)).isInstanceOf(ConsentDeniedException.class) + .hasMessageContaining("User denied consent for tool 'testTool' execution"); + + verify(delegate, never()).call(any()); + } + + @Test + void callWithExistingValidConsent() { + // Given + Map parameters = Map.of("param1", "value1"); + Object expectedResult = "Tool executed successfully"; + + when(delegate.getName()).thenReturn("testTool"); + when(consentManager.hasValidConsent("testTool", ConsentLevel.EVERY_TIME, new String[0])).thenReturn(true); + when(delegate.call(parameters)).thenReturn(expectedResult); + + // When + Object result = consentAwareCallback.call(parameters); + + // Then + assertThat(result).isEqualTo(expectedResult); + verify(delegate).call(parameters); + verify(consentManager, never()).requestConsent(any(), any(), any(), any(), any()); + } + + @Test + void messageSubstitutionWithParameters() { + // Given + Map parameters = new HashMap<>(); + parameters.put("bookId", "12345"); + parameters.put("title", "Spring in Action"); + + when(requiresConsent.message()).thenReturn("Delete book {bookId} titled '{title}'?"); + consentAwareCallback = new ConsentAwareToolCallback(delegate, consentManager, requiresConsent); + + when(delegate.getName()).thenReturn("deleteBook"); + when(consentManager.hasValidConsent("deleteBook", ConsentLevel.EVERY_TIME, new String[0])).thenReturn(false); + when(consentManager.requestConsent(eq("deleteBook"), eq("Delete book 12345 titled 'Spring in Action'?"), any(), + any(), eq(parameters))) + .thenReturn(true); + when(delegate.call(parameters)).thenReturn("Deleted"); + + // When + Object result = consentAwareCallback.call(parameters); + + // Then + assertThat(result).isEqualTo("Deleted"); + verify(consentManager).requestConsent(eq("deleteBook"), eq("Delete book 12345 titled 'Spring in Action'?"), + any(), any(), eq(parameters)); + } + + @Test + void delegateMethodsAreProxied() { + // Given + when(delegate.getName()).thenReturn("testTool"); + when(delegate.getDescription()).thenReturn("Test tool description"); + ToolDefinition toolDef = mock(ToolDefinition.class); + when(delegate.getToolDefinition()).thenReturn(toolDef); + + // When/Then + assertThat(consentAwareCallback.getName()).isEqualTo("testTool"); + assertThat(consentAwareCallback.getDescription()).isEqualTo("Test tool description"); + assertThat(consentAwareCallback.getToolDefinition()).isEqualTo(toolDef); + } + + @Test + void gettersReturnCorrectValues() { + // When/Then + assertThat(consentAwareCallback.getDelegate()).isEqualTo(delegate); + assertThat(consentAwareCallback.getRequiresConsent()).isEqualTo(requiresConsent); + } + +} diff --git a/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/DefaultConsentManagerTests.java b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/DefaultConsentManagerTests.java new file mode 100644 index 00000000000..ae06431a6e6 --- /dev/null +++ b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/DefaultConsentManagerTests.java @@ -0,0 +1,263 @@ +/* + * 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.consent; + +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Tests for {@link DefaultConsentManager}. + * + * @author Hyunjoon Park + */ +class DefaultConsentManagerTests { + + private DefaultConsentManager consentManager; + + private AtomicInteger consentRequestCount; + + private BiFunction, Boolean> alwaysApproveHandler; + + private BiFunction, Boolean> alwaysDenyHandler; + + @BeforeEach + void setUp() { + consentRequestCount = new AtomicInteger(0); + alwaysApproveHandler = (message, params) -> { + consentRequestCount.incrementAndGet(); + return true; + }; + alwaysDenyHandler = (message, params) -> { + consentRequestCount.incrementAndGet(); + return false; + }; + } + + @Test + void everyTimeConsentAlwaysRequiresNewConsent() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - First call + boolean granted1 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.EVERY_TIME, new String[0], + params); + + // Then + assertThat(granted1).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Second call + boolean granted2 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.EVERY_TIME, new String[0], + params); + + // Then + assertThat(granted2).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(2); + } + + @Test + void sessionConsentRemembersApproval() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - First call + boolean granted1 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], + params); + + // Then + assertThat(granted1).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Second call (should use stored consent) + boolean granted2 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], + params); + + // Then + assertThat(granted2).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); // No new request + } + + @Test + void rememberConsentPersistsApproval() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - First call + boolean granted1 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.REMEMBER, new String[0], + params); + + // Then + assertThat(granted1).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Multiple subsequent calls + for (int i = 0; i < 5; i++) { + boolean granted = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.REMEMBER, new String[0], + params); + assertThat(granted).isTrue(); + } + + // Then - Still only one consent request + assertThat(consentRequestCount.get()).isEqualTo(1); + } + + @Test + void deniedConsentIsNotStored() { + // Given + consentManager = new DefaultConsentManager(alwaysDenyHandler); + Map params = Map.of("param", "value"); + + // When - First call (denied) + boolean granted1 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], + params); + + // Then + assertThat(granted1).isFalse(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Second call (should request again since previous was denied) + boolean granted2 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], + params); + + // Then + assertThat(granted2).isFalse(); + assertThat(consentRequestCount.get()).isEqualTo(2); + } + + @Test + void consentWithCategories() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + String[] categories = { "destructive", "data-modification" }; + + // When - Grant consent for specific categories + boolean granted1 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, categories, params); + + // Then + assertThat(granted1).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Same tool, same categories (should use stored consent) + boolean granted2 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, categories, params); + + // Then + assertThat(granted2).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Same tool, different categories (should request new consent) + String[] differentCategories = { "read-only" }; + boolean granted3 = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, differentCategories, + params); + + // Then + assertThat(granted3).isTrue(); + assertThat(consentRequestCount.get()).isEqualTo(2); + } + + @Test + void revokeConsentForSpecificTool() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - Grant consent + consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], params); + assertThat(consentRequestCount.get()).isEqualTo(1); + + // When - Revoke consent + consentManager.revokeConsent("tool1", new String[0]); + + // When - Request consent again (should need new approval) + consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], params); + + // Then + assertThat(consentRequestCount.get()).isEqualTo(2); + } + + @Test + void clearAllConsents() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - Grant consent for multiple tools + consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], params); + consentManager.requestConsent("tool2", "Approve?", ConsentLevel.REMEMBER, new String[0], params); + assertThat(consentRequestCount.get()).isEqualTo(2); + + // When - Clear all consents + consentManager.clearAllConsents(); + + // When - Request consent again (should need new approvals) + consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], params); + consentManager.requestConsent("tool2", "Approve?", ConsentLevel.REMEMBER, new String[0], params); + + // Then + assertThat(consentRequestCount.get()).isEqualTo(4); + } + + @Test + void defaultConstructorAlwaysDeniesConsent() { + // Given + consentManager = new DefaultConsentManager(); + Map params = Map.of("param", "value"); + + // When + boolean granted = consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], + params); + + // Then + assertThat(granted).isFalse(); + } + + @Test + void hasValidConsentChecksStoredConsent() { + // Given + consentManager = new DefaultConsentManager(alwaysApproveHandler); + Map params = Map.of("param", "value"); + + // When - No consent stored yet + boolean hasConsent1 = consentManager.hasValidConsent("tool1", ConsentLevel.SESSION, new String[0]); + + // Then + assertThat(hasConsent1).isFalse(); + + // When - Grant consent + consentManager.requestConsent("tool1", "Approve?", ConsentLevel.SESSION, new String[0], params); + + // Then - Now has valid consent + boolean hasConsent2 = consentManager.hasValidConsent("tool1", ConsentLevel.SESSION, new String[0]); + assertThat(hasConsent2).isTrue(); + + // When - Check EVERY_TIME consent (always false) + boolean hasConsent3 = consentManager.hasValidConsent("tool1", ConsentLevel.EVERY_TIME, new String[0]); + assertThat(hasConsent3).isFalse(); + } + +} diff --git a/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/RequiresConsentIntegrationTests.java b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/RequiresConsentIntegrationTests.java new file mode 100644 index 00000000000..27542ebac00 --- /dev/null +++ b/spring-ai-model/src/test/java/org/springframework/ai/tool/consent/RequiresConsentIntegrationTests.java @@ -0,0 +1,125 @@ +/* + * 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.consent; + +import java.util.ArrayList; +import java.util.List; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import org.springframework.ai.tool.annotation.RequiresConsent; +import org.springframework.ai.tool.annotation.RequiresConsent.ConsentLevel; +import org.springframework.ai.tool.annotation.Tool; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Integration tests demonstrating the use of @RequiresConsent with @Tool annotations. + * + * @author Hyunjoon Park + */ +class RequiresConsentIntegrationTests { + + private List consentRequests; + + private ConsentManager consentManager; + + private BookService bookService; + + @BeforeEach + void setUp() { + consentRequests = new ArrayList<>(); + + // Create a consent manager that logs requests and approves based on message + // content + consentManager = new DefaultConsentManager((message, params) -> { + consentRequests.add(message); + // Approve if message contains "approve" (case insensitive) + return message.toLowerCase().contains("approve"); + }); + + bookService = new BookService(); + } + + @Test + void demonstrateRequiresConsentUsage() { + // This test demonstrates how the @RequiresConsent annotation would be used + // in practice with tool methods + + // The deleteBook method requires consent every time + assertThat(bookService.getClass().getDeclaredMethods()) + .filteredOn(method -> method.getName().equals("deleteBook")) + .hasSize(1) + .allSatisfy(method -> { + assertThat(method.isAnnotationPresent(Tool.class)).isTrue(); + assertThat(method.isAnnotationPresent(RequiresConsent.class)).isTrue(); + + RequiresConsent consent = method.getAnnotation(RequiresConsent.class); + assertThat(consent.message()).contains("{bookId}"); + assertThat(consent.level()).isEqualTo(ConsentLevel.EVERY_TIME); + }); + + // The updateBookPrice method requires session-level consent + assertThat(bookService.getClass().getDeclaredMethods()) + .filteredOn(method -> method.getName().equals("updateBookPrice")) + .hasSize(1) + .allSatisfy(method -> { + assertThat(method.isAnnotationPresent(Tool.class)).isTrue(); + assertThat(method.isAnnotationPresent(RequiresConsent.class)).isTrue(); + + RequiresConsent consent = method.getAnnotation(RequiresConsent.class); + assertThat(consent.level()).isEqualTo(ConsentLevel.SESSION); + assertThat(consent.categories()).contains("financial"); + }); + } + + /** + * Example service class demonstrating @RequiresConsent usage. + */ + static class BookService { + + @Tool(description = "Deletes a book from the database") + @RequiresConsent(message = "The book with ID {bookId} will be permanently deleted. Do you approve?", + level = ConsentLevel.EVERY_TIME) + public String deleteBook(String bookId) { + return "Book " + bookId + " deleted"; + } + + @Tool(description = "Updates the price of a book") + @RequiresConsent(message = "Update book {bookId} price from ${oldPrice} to ${newPrice}? Please approve.", + level = ConsentLevel.SESSION, categories = { "financial", "data-modification" }) + public String updateBookPrice(String bookId, double oldPrice, double newPrice) { + return String.format("Book %s price updated from %.2f to %.2f", bookId, oldPrice, newPrice); + } + + @Tool(description = "Gets book information") + // No @RequiresConsent - this operation doesn't need consent + public String getBook(String bookId) { + return "Book details for " + bookId; + } + + @Tool(description = "Sends book recommendation email") + @RequiresConsent(message = "Send book recommendations to {email}? (We'll remember your preference)", + level = ConsentLevel.REMEMBER, categories = { "communication", "marketing" }) + public String sendBookRecommendations(String email) { + return "Recommendations sent to " + email; + } + + } + +}