opensource/mateclaw
2
0
Fork 0
mirror of https://gitee.com/mateos/mateclaw.git synced 2026-09-13 03:13:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix(agent): cache runtime tool name lookups

Browse Source
This commit is contained in:
matevip 2026-08-14 04:50:13 -04:00
parent 50c6eb92f9
commit f629b04f2a
3 changed files with 65 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mateclaw-server/src/main/java/vip/mate/agent/AgentToolSet.java
View File

@ -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<String> allNames() {
return aliasIndex.keySet();
}
/**
* 获取原始的 @Tool Bean 列表
*/

97
mateclaw-server/src/main/java/vip/mate/tool/ToolRegistry.java
View File

@ -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<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) {}
@ -58,6 +64,7 @@ public class ToolRegistry {
*/
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());
}
@ -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<String> availableFunctionNames() {
Set<String> names = new java.util.HashSet<>();
Set<String> disabledBeanNames = toolMapper.selectList(
new LambdaQueryWrapper<ToolEntity>()
.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<String, Object> beans = applicationContext.getBeansWithAnnotation(Component.class);
for (Map.Entry<String, Object> 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<String, ToolCallbackProvider> 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();
}
/**

5
mateclaw-server/src/main/java/vip/mate/tool/service/ToolService.java
View File

@ -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);
}