mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-14 03:33:43 +08:00
396 lines
17 KiB
Java
396 lines
17 KiB
Java
package vip.mate.tool;
|
||
|
||
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
|
||
import lombok.RequiredArgsConstructor;
|
||
import lombok.extern.slf4j.Slf4j;
|
||
import org.springframework.ai.tool.annotation.Tool;
|
||
import org.springframework.boot.context.event.ApplicationReadyEvent;
|
||
import org.springframework.context.ApplicationContext;
|
||
import org.springframework.context.event.EventListener;
|
||
import org.springframework.stereotype.Component;
|
||
import vip.mate.tool.model.ToolEntity;
|
||
import vip.mate.tool.repository.ToolMapper;
|
||
|
||
import org.springframework.ai.support.ToolCallbacks;
|
||
import org.springframework.ai.tool.ToolCallback;
|
||
import org.springframework.ai.tool.ToolCallbackProvider;
|
||
import vip.mate.agent.AgentToolSet;
|
||
import vip.mate.i18n.I18nService;
|
||
import vip.mate.i18n.LocaleAwareToolCallback;
|
||
|
||
import java.util.ArrayList;
|
||
import java.util.Collections;
|
||
import java.util.IdentityHashMap;
|
||
import java.util.LinkedHashMap;
|
||
import java.util.LinkedHashSet;
|
||
import java.util.List;
|
||
import java.util.Map;
|
||
import java.util.Set;
|
||
import java.util.concurrent.CompletableFuture;
|
||
import java.util.concurrent.CopyOnWriteArrayList;
|
||
import java.util.function.Supplier;
|
||
import java.util.stream.Collectors;
|
||
|
||
/**
|
||
* 工具注册中心
|
||
* 管理所有可供 Agent 使用的工具（内置 + 自定义）
|
||
* 工具启用状态由数据库 mate_tool 表的 enabled 字段控制
|
||
*
|
||
* @author MateClaw Team
|
||
*/
|
||
@Slf4j
|
||
@Component
|
||
@RequiredArgsConstructor
|
||
public class ToolRegistry {
|
||
|
||
private final ApplicationContext applicationContext;
|
||
private final ToolMapper toolMapper;
|
||
private final I18nService i18nService;
|
||
|
||
// ==================== Plugin Tools ====================
|
||
|
||
/** Plugin-registered tool entries with lazy availability checks */
|
||
private final CopyOnWriteArrayList<PluginToolEntry> pluginTools = new CopyOnWriteArrayList<>();
|
||
|
||
private final Object enabledToolSetLock = new Object();
|
||
private volatile AgentToolSet enabledToolSetCache;
|
||
|
||
/** A tool entry registered by a plugin */
|
||
public record PluginToolEntry(ToolCallback callback, Supplier<Boolean> availabilityCheck) {}
|
||
|
||
/**
|
||
* Register a tool from a plugin with an availability check.
|
||
* The check is evaluated lazily each time the tool set is built.
|
||
*/
|
||
public void registerPluginTool(ToolCallback callback, Supplier<Boolean> availabilityCheck) {
|
||
pluginTools.add(new PluginToolEntry(callback, availabilityCheck != null ? availabilityCheck : () -> true));
|
||
invalidateEnabledToolSetCache("plugin-tool-registered:" + callback.getToolDefinition().name());
|
||
log.info("Plugin tool registered: {}", callback.getToolDefinition().name());
|
||
}
|
||
|
||
/**
|
||
* Unregister a plugin tool by name.
|
||
*/
|
||
public void unregisterPluginTool(String toolName) {
|
||
pluginTools.removeIf(entry -> entry.callback().getToolDefinition().name().equals(toolName));
|
||
invalidateEnabledToolSetCache("plugin-tool-unregistered:" + toolName);
|
||
log.info("Plugin tool unregistered: {}", toolName);
|
||
}
|
||
|
||
/**
|
||
* Snapshot plugin callbacks that are currently available to the runtime.
|
||
* Used by the agent tool picker so plugin tools can be bound per agent
|
||
* through the same {@code mate_agent_tool.tool_name} path as other tools.
|
||
*/
|
||
public List<ToolCallback> listAvailablePluginTools() {
|
||
List<ToolCallback> out = new ArrayList<>();
|
||
for (PluginToolEntry entry : pluginTools) {
|
||
try {
|
||
if (entry.callback() != null && Boolean.TRUE.equals(entry.availabilityCheck().get())) {
|
||
out.add(entry.callback());
|
||
}
|
||
} catch (Exception e) {
|
||
String name = entry.callback() != null && entry.callback().getToolDefinition() != null
|
||
? entry.callback().getToolDefinition().name()
|
||
: "<unknown>";
|
||
log.warn("Plugin tool availability check failed for {}: {}", name, e.getMessage());
|
||
}
|
||
}
|
||
return List.copyOf(out);
|
||
}
|
||
|
||
public void invalidateEnabledToolSetCache(String reason) {
|
||
enabledToolSetCache = null;
|
||
log.debug("Enabled AgentToolSet cache invalidated: {}", reason);
|
||
}
|
||
|
||
@EventListener(ApplicationReadyEvent.class)
|
||
public void prewarmEnabledToolSetCache() {
|
||
CompletableFuture.runAsync(() -> {
|
||
try {
|
||
getEnabledToolSet();
|
||
} catch (Exception e) {
|
||
log.debug("Enabled AgentToolSet prewarm skipped: {}", e.getMessage());
|
||
}
|
||
});
|
||
}
|
||
|
||
@EventListener
|
||
public void onMcpServerChanged(vip.mate.tool.mcp.event.McpServerChangedEvent event) {
|
||
invalidateEnabledToolSetCache("mcp-server-changed:" + event.reason());
|
||
prewarmEnabledToolSetCache();
|
||
}
|
||
|
||
@EventListener
|
||
public void onMcpConnectionLost(vip.mate.tool.mcp.event.McpConnectionLostEvent event) {
|
||
invalidateEnabledToolSetCache("mcp-connection-lost:" + event.serverId());
|
||
}
|
||
|
||
/**
|
||
* 获取所有已启用的工具 Bean（Spring AI @Tool 注解方式）
|
||
* 通过数据库 enabled 标志过滤，确保 UI 开关真正生效
|
||
*/
|
||
public List<Object> getEnabledTools() {
|
||
return List.copyOf(getEnabledToolBeansByName().values());
|
||
}
|
||
|
||
/**
|
||
* Iterate Spring beans once, returning a {@code beanName → bean} map of every
|
||
* currently-enabled @Tool bean.
|
||
* <p>
|
||
* 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<String, Object> getEnabledToolBeansByName() {
|
||
return getToolBeansByName(true);
|
||
}
|
||
|
||
/**
|
||
* Iterate Spring beans once, returning a {@code beanName → bean} map of
|
||
* every @Tool bean. When {@code enabledOnly} is true, DB rows with
|
||
* {@code enabled=false} are excluded; when false, disabled rows are kept for
|
||
* admin metadata use cases where the UI still needs to resolve aliases.
|
||
*/
|
||
private LinkedHashMap<String, Object> getToolBeansByName(boolean enabledOnly) {
|
||
// 1. 从数据库获取明确禁用的 beanName 黑名单
|
||
// 逻辑：只有 DB 中存在记录且 enabled=false 的才跳过
|
||
// DB 中没有记录的 bean 默认启用（向后兼容 + 新工具自动可用）
|
||
Set<String> disabledBeanNames = enabledOnly
|
||
? toolMapper.selectList(
|
||
new LambdaQueryWrapper<ToolEntity>()
|
||
.eq(ToolEntity::getEnabled, false)
|
||
.isNotNull(ToolEntity::getBeanName)
|
||
).stream()
|
||
.map(ToolEntity::getBeanName)
|
||
.collect(Collectors.toSet())
|
||
: Set.of();
|
||
|
||
LinkedHashMap<String, Object> enabled = new LinkedHashMap<>();
|
||
|
||
// 2. 扫描 Spring 容器中所有带 @Tool 方法的 Bean
|
||
Map<String, Object> beans = applicationContext.getBeansWithAnnotation(Component.class);
|
||
for (Map.Entry<String, Object> entry : beans.entrySet()) {
|
||
String beanName = entry.getKey();
|
||
Object bean = entry.getValue();
|
||
|
||
boolean hasToolMethod = java.util.Arrays.stream(bean.getClass().getMethods())
|
||
.anyMatch(m -> m.isAnnotationPresent(Tool.class));
|
||
|
||
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 {} tools: {}", enabledOnly ? "enabled" : "registered", enabled.size());
|
||
return enabled;
|
||
}
|
||
|
||
/**
|
||
* Build an alias index for every registered {@code @Tool} bean, including
|
||
* rows disabled in DB. This is for admin display only; runtime tool
|
||
* exposure must continue to use {@link #getEnabledToolSet()}.
|
||
*/
|
||
public AgentToolSet getAllToolBeanSetForAdmin() {
|
||
LinkedHashMap<String, Object> beansByName = getToolBeansByName(false);
|
||
List<Object> toolBeans = new ArrayList<>(beansByName.values());
|
||
IdentityHashMap<Object, String> nameByBean = new IdentityHashMap<>();
|
||
for (Map.Entry<String, Object> e : beansByName.entrySet()) {
|
||
nameByBean.put(e.getValue(), e.getKey());
|
||
}
|
||
List<ToolCallback> callbacks = new ArrayList<>();
|
||
for (Object bean : toolBeans) {
|
||
Collections.addAll(callbacks, ToolCallbacks.from(bean));
|
||
}
|
||
return AgentToolSet.fromCallbacks(toolBeans, callbacks, nameByBean::get);
|
||
}
|
||
|
||
/**
|
||
* Resolve aliases for currently enabled built-in {@code @Tool} beans without
|
||
* touching {@link ToolCallbackProvider}s. This is intentionally narrower than
|
||
* {@link #getEnabledToolSet()}: disclosure-tier snapshots only need to bridge
|
||
* {@code mate_tool.name}/{@code bean_name} onto built-in function names, and
|
||
* calling providers here would synchronously enumerate MCP tools on the chat
|
||
* hot path.
|
||
*/
|
||
public Set<String> enabledToolBeanFunctionNamesFor(Set<String> aliases) {
|
||
if (aliases == null || aliases.isEmpty()) {
|
||
return Set.of();
|
||
}
|
||
Map<String, Set<String>> index = enabledToolBeanFunctionNameIndex();
|
||
LinkedHashSet<String> out = new LinkedHashSet<>();
|
||
for (String alias : aliases) {
|
||
Set<String> hits = index.get(alias);
|
||
if (hits != null) {
|
||
out.addAll(hits);
|
||
}
|
||
}
|
||
return out;
|
||
}
|
||
|
||
/**
|
||
* Build {@code alias -> @Tool function names} for enabled built-in tool beans.
|
||
* The aliases mirror {@link AgentToolSet}: function name, Spring bean name,
|
||
* and Java simple class name. Provider/MCP callbacks are deliberately absent.
|
||
*/
|
||
public Map<String, Set<String>> enabledToolBeanFunctionNameIndex() {
|
||
LinkedHashMap<String, Object> beansByName = getEnabledToolBeansByName();
|
||
Map<String, Set<String>> index = new LinkedHashMap<>();
|
||
for (Map.Entry<String, Object> entry : beansByName.entrySet()) {
|
||
String beanName = entry.getKey();
|
||
Object bean = entry.getValue();
|
||
ToolCallback[] callbacks = ToolCallbacks.from(bean);
|
||
LinkedHashSet<String> functionNames = new LinkedHashSet<>();
|
||
for (ToolCallback cb : callbacks) {
|
||
if (cb != null && cb.getToolDefinition() != null) {
|
||
functionNames.add(cb.getToolDefinition().name());
|
||
}
|
||
}
|
||
if (functionNames.isEmpty()) {
|
||
continue;
|
||
}
|
||
putAlias(index, beanName, functionNames);
|
||
putAlias(index, bean.getClass().getSimpleName(), functionNames);
|
||
for (String functionName : functionNames) {
|
||
putAlias(index, functionName, Set.of(functionName));
|
||
}
|
||
}
|
||
return index;
|
||
}
|
||
|
||
private static void putAlias(Map<String, Set<String>> index, String alias, Set<String> functionNames) {
|
||
if (alias == null || alias.isBlank() || functionNames == null || functionNames.isEmpty()) {
|
||
return;
|
||
}
|
||
index.computeIfAbsent(alias, ignored -> new LinkedHashSet<>()).addAll(functionNames);
|
||
}
|
||
|
||
/**
|
||
* 获取统一的 AgentToolSet（包含 @Tool Bean + ToolCallbackProvider）
|
||
* <p>
|
||
* 同时收集：
|
||
* 1. 当前启用的 @Tool bean
|
||
* 2. 当前容器中所有 ToolCallbackProvider（MCP server 等）
|
||
*/
|
||
public AgentToolSet getEnabledToolSet() {
|
||
AgentToolSet cached = enabledToolSetCache;
|
||
if (cached != null) {
|
||
return cached;
|
||
}
|
||
synchronized (enabledToolSetLock) {
|
||
cached = enabledToolSetCache;
|
||
if (cached != null) {
|
||
return cached;
|
||
}
|
||
AgentToolSet built = buildEnabledToolSet();
|
||
enabledToolSetCache = built;
|
||
return built;
|
||
}
|
||
}
|
||
|
||
private AgentToolSet buildEnabledToolSet() {
|
||
// 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<String, Object> beansByName = getEnabledToolBeansByName();
|
||
List<Object> toolBeans = new ArrayList<>(beansByName.values());
|
||
IdentityHashMap<Object, String> nameByBean = new IdentityHashMap<>();
|
||
for (Map.Entry<String, Object> e : beansByName.entrySet()) {
|
||
nameByBean.put(e.getValue(), e.getKey());
|
||
}
|
||
|
||
Map<String, ToolCallbackProvider> providerBeans = applicationContext.getBeansOfType(ToolCallbackProvider.class);
|
||
List<ToolCallbackProvider> providers = new ArrayList<>(providerBeans.values());
|
||
|
||
// 对内置工具 callback 应用 i18n 描述包装
|
||
List<ToolCallback> localizedCallbacks = new ArrayList<>();
|
||
for (Object bean : toolBeans) {
|
||
ToolCallback[] cbs = ToolCallbacks.from(bean);
|
||
for (ToolCallback cb : cbs) {
|
||
String toolName = cb.getToolDefinition().name();
|
||
String descKey = "tool." + toolName + ".desc";
|
||
// The i18n description is an optional override: tools without a
|
||
// bundle entry (e.g. wiki tools) keep the description declared on
|
||
// their @Tool annotation. Use msgOptional so an absent key is not
|
||
// logged as a "missing key" — that is expected, not a fault.
|
||
String localizedDesc = i18nService.msgOptional(descKey);
|
||
if (localizedDesc != null) {
|
||
localizedCallbacks.add(new LocaleAwareToolCallback(cb, localizedDesc));
|
||
} else {
|
||
localizedCallbacks.add(cb);
|
||
}
|
||
}
|
||
}
|
||
|
||
// MCP provider callbacks 不做 i18n 包装（MCP 工具自行管理描述）
|
||
for (ToolCallbackProvider provider : providers) {
|
||
ToolCallback[] cbs = provider.getToolCallbacks();
|
||
if (cbs != null) {
|
||
Collections.addAll(localizedCallbacks, cbs);
|
||
}
|
||
}
|
||
|
||
// Plugin tool callbacks — evaluate availability checks lazily
|
||
int pluginToolCount = 0;
|
||
for (PluginToolEntry entry : pluginTools) {
|
||
try {
|
||
if (Boolean.TRUE.equals(entry.availabilityCheck().get())) {
|
||
localizedCallbacks.add(entry.callback());
|
||
pluginToolCount++;
|
||
} else {
|
||
log.debug("Plugin tool excluded (availability check failed): {}",
|
||
entry.callback().getToolDefinition().name());
|
||
}
|
||
} catch (Exception e) {
|
||
log.warn("Plugin tool availability check failed for {}: {}",
|
||
entry.callback().getToolDefinition().name(), e.getMessage());
|
||
}
|
||
}
|
||
|
||
log.info("Building AgentToolSet: toolBeans={}, providers={}, pluginTools={}, totalCallbacks={}",
|
||
toolBeans.size(), providers.size(), pluginToolCount, localizedCallbacks.size());
|
||
return AgentToolSet.fromCallbacks(toolBeans, localizedCallbacks, nameByBean::get);
|
||
}
|
||
|
||
/**
|
||
* Returns every runtime identifier by which a currently-enabled tool can be
|
||
* referenced — SKILL.md authors use all three conventions interchangeably:
|
||
* <ul>
|
||
* <li>{@code @Tool} function name (e.g. {@code browser_use}, {@code runSkillScript})</li>
|
||
* <li>Spring bean name (e.g. {@code browserUseTool}, {@code skillScriptTool})</li>
|
||
* <li>MCP tool id / plugin tool name (routed via {@code ToolCallbackProvider})</li>
|
||
* </ul>
|
||
* Returning the union lets {@link vip.mate.skill.runtime.SkillDependencyChecker}
|
||
* accept whichever convention a skill happens to declare.
|
||
*/
|
||
public Set<String> availableFunctionNames() {
|
||
return getEnabledToolSet().allNames();
|
||
}
|
||
|
||
/**
|
||
* 获取数据库中的工具配置列表（全部）
|
||
*/
|
||
public List<ToolEntity> listToolEntities() {
|
||
return toolMapper.selectList(new LambdaQueryWrapper<ToolEntity>()
|
||
.orderByDesc(ToolEntity::getBuiltin)
|
||
.orderByAsc(ToolEntity::getName));
|
||
}
|
||
|
||
/**
|
||
* 获取已启用的工具配置列表
|
||
*/
|
||
public List<ToolEntity> listEnabledToolEntities() {
|
||
return toolMapper.selectList(new LambdaQueryWrapper<ToolEntity>()
|
||
.eq(ToolEntity::getEnabled, true)
|
||
.orderByAsc(ToolEntity::getName));
|
||
}
|
||
}
|