diff --git a/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java b/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java index a15d2472..1f606ebe 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java @@ -5,8 +5,8 @@ import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.ToolCallbackProvider; import java.util.*; +import java.util.function.Function; import java.util.stream.Collectors; -import java.util.LinkedHashMap; /** * Agent 统一工具集合 @@ -14,6 +14,20 @@ import java.util.LinkedHashMap; * 将 @Tool Bean、ToolCallbackProvider、MCP server 暴露的 tool callbacks * 统一收集为一致的 ToolCallback 列表,供 StateGraph 节点使用。 * + *

Alias index — why one tool has multiple names

+ * Each tool can be referenced by several equivalent identifiers: + * + * Filtering operations ({@link #withAllowedToolsOnly}, {@link #withDeniedToolsFiltered}, + * {@link #excluding}) accept any of these aliases, so callers don't need to know which + * naming convention the persistence layer happens to use. This is the same pattern Spring's + * {@code BeanFactory} uses for bean names + aliases. + * * @author MateClaw Team */ public class AgentToolSet { @@ -21,26 +35,69 @@ public class AgentToolSet { private final List toolBeans; private final List callbacks; private final Map callbackByName; + /** + * Alias → callbacks. One alias may resolve to multiple callbacks + * (e.g. a Spring bean name pointing at a class that exposes several {@code @Tool} methods), + * which is why values are sets. + */ + private final Map> aliasIndex; - private AgentToolSet(List toolBeans, List callbacks) { + private AgentToolSet(List toolBeans, List callbacks, + Function beanNameResolver) { this.toolBeans = List.copyOf(toolBeans); // 按工具名去重:内置工具在前(先添加),MCP 工具在后,同名时保留内置工具 // 使用 LinkedHashMap 保证插入顺序,确保内置工具始终排在 MCP 工具前面(影响 LLM 工具选择倾向) - this.callbackByName = callbacks.stream() + LinkedHashMap byName = callbacks.stream() .collect(Collectors.toMap( cb -> cb.getToolDefinition().name(), cb -> cb, (a, b) -> a, LinkedHashMap::new)); + this.callbackByName = byName; // callbacks 列表也使用去重后的结果,避免 Spring AI ToolCallingChatOptions 校验重名报错 - this.callbacks = List.copyOf(callbackByName.values()); + this.callbacks = List.copyOf(byName.values()); + this.aliasIndex = buildAliasIndex(this.toolBeans, byName, beanNameResolver); } + /** + * Internal constructor for {@link #rebuild} — preserves a pre-filtered alias index + * so we don't need {@code beanNameResolver} on every {@code with*} call. + */ + private AgentToolSet(List toolBeans, List callbacks, + Map> precomputedAliasIndex) { + this.toolBeans = List.copyOf(toolBeans); + LinkedHashMap byName = callbacks.stream() + .collect(Collectors.toMap( + cb -> cb.getToolDefinition().name(), + cb -> cb, + (a, b) -> a, + LinkedHashMap::new)); + this.callbackByName = byName; + this.callbacks = List.copyOf(byName.values()); + this.aliasIndex = Map.copyOf(precomputedAliasIndex); + } + + /** No-op resolver for callers that don't have access to Spring bean names. */ + private static final Function NO_BEAN_NAMES = bean -> null; + /** * 从预构建的 ToolCallback 列表构建工具集(用于 i18n 等需要包装 callback 的场景) */ public static AgentToolSet fromCallbacks(List toolBeans, List callbacks) { - return new AgentToolSet(toolBeans != null ? toolBeans : List.of(), callbacks); + return new AgentToolSet(toolBeans != null ? toolBeans : List.of(), callbacks, NO_BEAN_NAMES); + } + + /** + * Same as {@link #fromCallbacks(List, List)} but additionally indexes each tool bean by + * its Spring bean name and Java simple class name, so {@link #withAllowedToolsOnly} accepts + * any of those identifiers (in addition to the {@code @Tool} function name). + * + * @param beanNameResolver lookup from a tool bean instance to its Spring bean name; + * may return {@code null} if the bean has no registered name + */ + public static AgentToolSet fromCallbacks(List toolBeans, List callbacks, + Function beanNameResolver) { + return new AgentToolSet(toolBeans != null ? toolBeans : List.of(), callbacks, beanNameResolver); } /** @@ -67,36 +124,45 @@ public class AgentToolSet { } } - return new AgentToolSet(toolBeans != null ? toolBeans : List.of(), allCallbacks); + return new AgentToolSet(toolBeans != null ? toolBeans : List.of(), allCallbacks, NO_BEAN_NAMES); } /** * 过滤掉 denied 工具后返回新的 AgentToolSet。 * denied 工具不会暴露给模型,模型完全不知道它们的存在。 * - * @param deniedTools denied 工具名集合(为空或 null 时直接返回 this) + * @param deniedTools denied 工具名集合(接受 function name / bean name / class simple name; + * 为空或 null 时直接返回 this) */ public AgentToolSet withDeniedToolsFiltered(Set deniedTools) { if (deniedTools == null || deniedTools.isEmpty()) { return this; } - List filtered = new ArrayList<>(callbacks); - filtered.removeIf(cb -> deniedTools.contains(cb.getToolDefinition().name())); - return new AgentToolSet(toolBeans, filtered); + Set denied = resolveAliases(deniedTools); + if (denied.isEmpty()) { + return this; + } + List filtered = callbacks.stream() + .filter(cb -> !denied.contains(cb)) + .toList(); + return rebuild(filtered); } /** * 仅保留指定名称的工具(白名单模式,用于 per-agent 绑定) * - * @param allowedTools 允许的工具名集合(为 null 时直接返回 this,表示使用全局默认) + * @param allowedTools 允许的工具名集合(接受 function name / Spring bean name / Java class simple name; + * 为 null 时直接返回 this,表示使用全局默认) */ public AgentToolSet withAllowedToolsOnly(Set allowedTools) { if (allowedTools == null) { return this; // null = 无绑定,使用全局默认 } - List filtered = new ArrayList<>(callbacks); - filtered.removeIf(cb -> !allowedTools.contains(cb.getToolDefinition().name())); - return new AgentToolSet(toolBeans, filtered); + Set allowed = resolveAliases(allowedTools); + List filtered = callbacks.stream() + .filter(allowed::contains) + .toList(); + return rebuild(filtered); } /** @@ -122,15 +188,21 @@ public class AgentToolSet { /** * 返回排除指定工具名后的新 AgentToolSet + * + * @param toolNames 要排除的工具名集合(接受 function name / bean name / class simple name) */ public AgentToolSet excluding(Set toolNames) { if (toolNames == null || toolNames.isEmpty()) { return this; } + Set excluded = resolveAliases(toolNames); + if (excluded.isEmpty()) { + return this; + } List filtered = callbacks.stream() - .filter(cb -> !toolNames.contains(cb.getToolDefinition().name())) + .filter(cb -> !excluded.contains(cb)) .toList(); - return new AgentToolSet(toolBeans, filtered); + return rebuild(filtered); } /** @@ -146,4 +218,107 @@ public class AgentToolSet { public int size() { return callbacks.size(); } + + // ==================== Internals ==================== + + /** + * Resolve a set of aliases (any mix of function name / bean name / class simple name) + * into the set of {@link ToolCallback} instances they refer to. Unknown aliases are + * silently dropped — the caller is expected to be tolerant of stale persistence data. + */ + private Set resolveAliases(Set aliases) { + Set resolved = new LinkedHashSet<>(); + for (String alias : aliases) { + Set hits = aliasIndex.get(alias); + if (hits != null) { + resolved.addAll(hits); + } + } + return resolved; + } + + /** + * Reconstruct a new {@code AgentToolSet} after filtering callbacks, carrying forward + * only the alias entries whose targets survived. This avoids re-running + * {@link ToolCallbacks#from(Object)} reflection on every {@code with*} call. + */ + private AgentToolSet rebuild(List filteredCallbacks) { + Set survivors = new HashSet<>(filteredCallbacks); + Map> filteredAliases = new LinkedHashMap<>(); + for (Map.Entry> e : aliasIndex.entrySet()) { + Set kept = new LinkedHashSet<>(); + for (ToolCallback cb : e.getValue()) { + if (survivors.contains(cb)) { + kept.add(cb); + } + } + if (!kept.isEmpty()) { + filteredAliases.put(e.getKey(), Set.copyOf(kept)); + } + } + return new AgentToolSet(toolBeans, filteredCallbacks, filteredAliases); + } + + /** + * Build the alias index. Function names are always indexed (they are the runtime truth); + * bean names and class simple names are indexed when {@code beanNameResolver} is provided + * — typically only the production registry has the {@link org.springframework.context.ApplicationContext} + * needed to map bean instances to names. Unit tests that pass empty {@code toolBeans} + * naturally get a function-name-only index. + */ + private static Map> buildAliasIndex( + List toolBeans, + Map callbackByName, + Function beanNameResolver) { + + Map> aliases = new LinkedHashMap<>(); + + // 1. Always index by function name (the runtime identifier) + for (Map.Entry e : callbackByName.entrySet()) { + aliases.computeIfAbsent(e.getKey(), k -> new LinkedHashSet<>()).add(e.getValue()); + } + + // 2. If we have bean info, also index by Spring bean name and Java class simple name. + // A single bean may expose multiple @Tool methods → the alias maps to a set. + if (beanNameResolver != null) { + for (Object bean : toolBeans) { + String beanName = beanNameResolver.apply(bean); + String simpleName = bean.getClass().getSimpleName(); + + // Find which callbacks belong to this bean, looking them up in the + // (possibly i18n-wrapped) callbackByName so we point at the same + // instances the rest of the set uses. + Set beanCallbacks = new LinkedHashSet<>(); + ToolCallback[] rawCallbacks; + try { + rawCallbacks = ToolCallbacks.from(bean); + } catch (Exception ignored) { + // Defensive: a misbehaving bean shouldn't break the whole tool set + continue; + } + for (ToolCallback raw : rawCallbacks) { + ToolCallback wrapped = callbackByName.get(raw.getToolDefinition().name()); + if (wrapped != null) { + beanCallbacks.add(wrapped); + } + } + if (beanCallbacks.isEmpty()) { + continue; + } + if (beanName != null && !beanName.isBlank()) { + aliases.computeIfAbsent(beanName, k -> new LinkedHashSet<>()).addAll(beanCallbacks); + } + if (simpleName != null && !simpleName.isBlank()) { + aliases.computeIfAbsent(simpleName, k -> new LinkedHashSet<>()).addAll(beanCallbacks); + } + } + } + + // Freeze inner sets + Map> frozen = new LinkedHashMap<>(); + for (Map.Entry> e : aliases.entrySet()) { + frozen.put(e.getKey(), Set.copyOf(e.getValue())); + } + return Map.copyOf(frozen); + } } diff --git a/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java b/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java index 256a3086..8a736fe1 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java @@ -18,6 +18,8 @@ import vip.mate.i18n.LocaleAwareToolCallback; import java.util.ArrayList; import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -71,6 +73,19 @@ public class ToolRegistry { * 通过数据库 enabled 标志过滤,确保 UI 开关真正生效 */ public List getEnabledTools() { + return List.copyOf(getEnabledToolBeansByName().values()); + } + + /** + * Iterate Spring beans once, returning a {@code beanName → bean} map of every + * currently-enabled @Tool bean. + *

+ * This is the single source of truth for "which @Tool beans should the agent see"; both + * {@link #getEnabledTools()} and {@link #getEnabledToolSet()} build on it. Returning + * {@link LinkedHashMap} preserves the discovery order from {@code getBeansWithAnnotation}, + * which {@link AgentToolSet} relies on (built-in tools first, MCP tools second). + */ + private LinkedHashMap getEnabledToolBeansByName() { // 1. 从数据库获取明确禁用的 beanName 黑名单 // 逻辑:只有 DB 中存在记录且 enabled=false 的才跳过 // DB 中没有记录的 bean 默认启用(向后兼容 + 新工具自动可用) @@ -82,7 +97,7 @@ public class ToolRegistry { .map(ToolEntity::getBeanName) .collect(Collectors.toSet()); - List tools = new ArrayList<>(); + LinkedHashMap enabled = new LinkedHashMap<>(); // 2. 扫描 Spring 容器中所有带 @Tool 方法的 Bean Map beans = applicationContext.getBeansWithAnnotation(Component.class); @@ -93,19 +108,20 @@ public class ToolRegistry { boolean hasToolMethod = java.util.Arrays.stream(bean.getClass().getMethods()) .anyMatch(m -> m.isAnnotationPresent(Tool.class)); - if (hasToolMethod) { - // 3. 只有 DB 中明确 enabled=false 的才跳过,其余全部启用 - if (disabledBeanNames.contains(beanName)) { - log.debug("Skipped disabled tool bean: {} (beanName={})", bean.getClass().getSimpleName(), beanName); - } else { - tools.add(bean); - log.debug("Registered tool bean: {} (beanName={})", bean.getClass().getSimpleName(), beanName); - } + if (!hasToolMethod) { + continue; + } + // 3. 只有 DB 中明确 enabled=false 的才跳过,其余全部启用 + if (disabledBeanNames.contains(beanName)) { + log.debug("Skipped disabled tool bean: {} (beanName={})", bean.getClass().getSimpleName(), beanName); + } else { + enabled.put(beanName, bean); + log.debug("Registered tool bean: {} (beanName={})", bean.getClass().getSimpleName(), beanName); } } - log.info("Total enabled tools: {}", tools.size()); - return tools; + log.info("Total enabled tools: {}", enabled.size()); + return enabled; } /** @@ -116,7 +132,16 @@ public class ToolRegistry { * 2. 当前容器中所有 ToolCallbackProvider(MCP server 等) */ public AgentToolSet getEnabledToolSet() { - List toolBeans = getEnabledTools(); + // Build both the bean list and the identity-based name lookup in one pass — the + // latter lets AgentToolSet's alias index resolve a saved binding like + // "BrowserUseTool" or "browserUseTool" back to the same callback as "browser_use". + LinkedHashMap beansByName = getEnabledToolBeansByName(); + List toolBeans = new ArrayList<>(beansByName.values()); + IdentityHashMap nameByBean = new IdentityHashMap<>(); + for (Map.Entry e : beansByName.entrySet()) { + nameByBean.put(e.getValue(), e.getKey()); + } + Map providerBeans = applicationContext.getBeansOfType(ToolCallbackProvider.class); List providers = new ArrayList<>(providerBeans.values()); @@ -164,7 +189,7 @@ public class ToolRegistry { log.info("Building AgentToolSet: toolBeans={}, providers={}, pluginTools={}, totalCallbacks={}", toolBeans.size(), providers.size(), pluginToolCount, localizedCallbacks.size()); - return AgentToolSet.fromCallbacks(toolBeans, localizedCallbacks); + return AgentToolSet.fromCallbacks(toolBeans, localizedCallbacks, nameByBean::get); } /**