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 e247aa4b..60226fea 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java @@ -179,6 +179,14 @@ public class AgentToolSet { return callbackByName; } + /** + * Every runtime identifier this set can resolve: function names plus any + * Spring bean / Java class aliases captured when the set was built. + */ + public Set allNames() { + return aliasIndex.keySet(); + } + /** * 获取原始的 @Tool Bean 列表 */ 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 b76e533c..de320a1c 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java @@ -4,7 +4,9 @@ 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; @@ -24,6 +26,7 @@ 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; @@ -49,6 +52,9 @@ public class ToolRegistry { /** Plugin-registered tool entries with lazy availability checks */ private final CopyOnWriteArrayList 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 availabilityCheck) {} @@ -58,6 +64,7 @@ public class ToolRegistry { */ public void registerPluginTool(ToolCallback callback, Supplier 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()); } @@ -66,9 +73,37 @@ public class ToolRegistry { */ 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); } + 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 开关真正生效 @@ -224,6 +259,22 @@ public class ToolRegistry { * 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". @@ -299,51 +350,7 @@ public class ToolRegistry { * accept whichever convention a skill happens to declare. */ public Set availableFunctionNames() { - Set names = new java.util.HashSet<>(); - - Set disabledBeanNames = toolMapper.selectList( - new LambdaQueryWrapper() - .eq(ToolEntity::getEnabled, false) - .isNotNull(ToolEntity::getBeanName) - ).stream().map(ToolEntity::getBeanName).collect(Collectors.toSet()); - - // 1. @Tool beans — register both the bean name and every function name exposed. - Map beans = applicationContext.getBeansWithAnnotation(Component.class); - for (Map.Entry entry : beans.entrySet()) { - String beanName = entry.getKey(); - Object bean = entry.getValue(); - if (disabledBeanNames.contains(beanName)) continue; - boolean hasToolMethod = java.util.Arrays.stream(bean.getClass().getMethods()) - .anyMatch(m -> m.isAnnotationPresent(Tool.class)); - if (!hasToolMethod) continue; - names.add(beanName); - for (ToolCallback cb : ToolCallbacks.from(bean)) { - names.add(cb.getToolDefinition().name()); - } - } - - // 2. MCP providers — only function names exist here. - Map providers = applicationContext.getBeansOfType(ToolCallbackProvider.class); - for (ToolCallbackProvider provider : providers.values()) { - ToolCallback[] cbs = provider.getToolCallbacks(); - if (cbs == null) continue; - for (ToolCallback cb : cbs) { - names.add(cb.getToolDefinition().name()); - } - } - - // 3. Plugin-registered tools — evaluate availability lazily so disabled plugins drop out. - for (PluginToolEntry entry : pluginTools) { - try { - if (Boolean.TRUE.equals(entry.availabilityCheck().get())) { - names.add(entry.callback().getToolDefinition().name()); - } - } catch (Exception ignored) { - // Unreachable plugin tools don't contribute to the set. - } - } - - return names; + return getEnabledToolSet().allNames(); } /** diff --git a/mateclaw-server/src/main/java/vip/mate/tool/service/ToolService.java b/mateclaw-server/src/main/java/vip/mate/tool/service/ToolService.java index fc52751d..1acabb9b 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/service/ToolService.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/service/ToolService.java @@ -47,6 +47,7 @@ public class ToolService { tool.setEnabled(true); } toolMapper.insert(tool); + toolRegistry.invalidateEnabledToolSetCache("tool-created:" + tool.getName()); return enrichRuntimeNames(tool); } @@ -55,9 +56,11 @@ public class ToolService { if (Boolean.TRUE.equals(existing.getBuiltin())) { existing.setEnabled(tool.getEnabled()); toolMapper.updateById(existing); + toolRegistry.invalidateEnabledToolSetCache("builtin-tool-updated:" + existing.getName()); return enrichRuntimeNames(existing); } toolMapper.updateById(tool); + toolRegistry.invalidateEnabledToolSetCache("tool-updated:" + tool.getName()); return enrichRuntimeNames(tool); } @@ -67,12 +70,14 @@ public class ToolService { throw new MateClawException("err.tool.builtin_readonly", "内置工具不可删除"); } toolMapper.deleteById(id); + toolRegistry.invalidateEnabledToolSetCache("tool-deleted:" + tool.getName()); } public ToolEntity toggleTool(Long id, boolean enabled) { ToolEntity tool = getTool(id); tool.setEnabled(enabled); toolMapper.updateById(tool); + toolRegistry.invalidateEnabledToolSetCache("tool-toggled:" + tool.getName()); return enrichRuntimeNames(tool); }