mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 11:13:43 +08:00
392 lines
17 KiB
Java
392 lines
17 KiB
Java
package vip.mate.tool.disclosure;
|
|
|
|
import cn.hutool.json.JSONArray;
|
|
import cn.hutool.json.JSONObject;
|
|
import cn.hutool.json.JSONUtil;
|
|
import lombok.RequiredArgsConstructor;
|
|
import lombok.extern.slf4j.Slf4j;
|
|
import org.springframework.ai.tool.ToolCallback;
|
|
import org.springframework.beans.factory.annotation.Value;
|
|
import org.springframework.stereotype.Service;
|
|
import vip.mate.agent.AgentToolSet;
|
|
import vip.mate.agent.context.TokenEstimator;
|
|
import vip.mate.tool.ToolRegistry;
|
|
import vip.mate.tool.mcp.model.McpServerEntity;
|
|
import vip.mate.tool.mcp.runtime.McpHashCollisionDetector;
|
|
import vip.mate.tool.mcp.service.McpServerService;
|
|
import vip.mate.tool.model.ToolEntity;
|
|
import vip.mate.tool.service.ToolService;
|
|
|
|
import java.util.ArrayList;
|
|
import java.util.Comparator;
|
|
import java.util.LinkedHashMap;
|
|
import java.util.LinkedHashSet;
|
|
import java.util.List;
|
|
import java.util.Map;
|
|
import java.util.Set;
|
|
|
|
/**
|
|
* Default tier resolver. Tier data is read from {@code mate_tool} and
|
|
* {@code mate_mcp_server} and cached in a short-lived snapshot so the per-turn
|
|
* {@link #split} does not hit the DB on every reasoning step. The PATCH
|
|
* endpoints call {@link #invalidate()} after changing a tier.
|
|
*/
|
|
@Slf4j
|
|
@Service
|
|
@RequiredArgsConstructor
|
|
public class DefaultToolDisclosureService implements ToolDisclosureService {
|
|
|
|
private static final long CACHE_TTL_MS = 30_000L;
|
|
|
|
/**
|
|
* Meta-tools that must always stay core: hiding them would make progressive
|
|
* disclosure unrecoverable (the model could neither search/call a deferred
|
|
* tool nor load a skill).
|
|
*/
|
|
private static final Set<String> ALWAYS_CORE = Set.of(
|
|
"enable_tool", "load_skill", "tool_search", "tool_describe", "tool_call");
|
|
|
|
/**
|
|
* Code-level extension defaults for builtin tools that may not yet have a
|
|
* {@code mate_tool} row when first resolved. A persisted
|
|
* {@code disclosure_tier} always overrides these.
|
|
*/
|
|
private static final Set<String> BUILTIN_EXTENSION_DEFAULTS = Set.of(
|
|
"image_generate", "music_generate", "video_generate", "model3d_generate", "browser_use");
|
|
|
|
private final ToolService toolService;
|
|
private final McpServerService mcpServerService;
|
|
private final ToolRegistry toolRegistry;
|
|
private final ToolUsageRecencyTracker usageRecencyTracker;
|
|
|
|
@Value("${mateclaw.tools.disclosure.mode:progressive}")
|
|
private String disclosureMode;
|
|
|
|
private volatile Snapshot snapshot;
|
|
|
|
private boolean legacyMode() {
|
|
return "legacy".equalsIgnoreCase(disclosureMode);
|
|
}
|
|
|
|
@Override
|
|
public DisclosureTier resolveTier(ToolCallback callback) {
|
|
if (callback == null || callback.getToolDefinition() == null) {
|
|
return DisclosureTier.CORE;
|
|
}
|
|
return resolveTierByName(callback.getToolDefinition().name());
|
|
}
|
|
|
|
@Override
|
|
public DisclosureTier resolveTierByName(String toolName) {
|
|
if (toolName == null || toolName.isBlank() || legacyMode()) {
|
|
return DisclosureTier.CORE;
|
|
}
|
|
if (ALWAYS_CORE.contains(toolName)) {
|
|
return DisclosureTier.CORE;
|
|
}
|
|
Snapshot snap = snapshot();
|
|
DisclosureTier dbTier = snap.builtinTierByName.get(toolName);
|
|
if (dbTier != null) {
|
|
return dbTier;
|
|
}
|
|
if (BUILTIN_EXTENSION_DEFAULTS.contains(toolName)) {
|
|
return DisclosureTier.EXTENSION;
|
|
}
|
|
Long serverId = snap.mcpToolToServerId.get(toolName);
|
|
if (serverId != null) {
|
|
// Move 5: MCP tools default to EXTENSION (on-demand exposure).
|
|
// A server with 20 tools flooding the CORE tool list makes it
|
|
// harder for the model to find the right builtin tool, and the
|
|
// MCP schemas are typically the heaviest part of the prompt.
|
|
// Users who want a server's tools visible by default can set
|
|
// disclosure_tier = core on the mate_mcp_server row.
|
|
return snap.serverTierById.getOrDefault(serverId, DisclosureTier.EXTENSION);
|
|
}
|
|
// Unknown source (ACP / dynamic-skill / plugin) — keep visible.
|
|
return DisclosureTier.CORE;
|
|
}
|
|
|
|
@Override
|
|
public ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions) {
|
|
return split(baseSet, enabledExtensions, Set.of());
|
|
}
|
|
|
|
@Override
|
|
public ToolDisclosureSplit split(AgentToolSet baseSet, Set<String> enabledExtensions,
|
|
Set<String> autoDemoted) {
|
|
List<ToolCallback> all = baseSet == null ? List.of() : baseSet.callbacks();
|
|
if (legacyMode()) {
|
|
return new ToolDisclosureSplit(all, List.of());
|
|
}
|
|
Set<String> enabled = enabledExtensions == null ? Set.of() : enabledExtensions;
|
|
Set<String> demoted = autoDemoted == null ? Set.of() : autoDemoted;
|
|
List<ToolCallback> active = new ArrayList<>(all.size());
|
|
List<ToolCallback> extensionCatalog = new ArrayList<>();
|
|
for (ToolCallback cb : all) {
|
|
String name = cb.getToolDefinition().name();
|
|
if (resolveTier(cb) == DisclosureTier.EXTENSION || demoted.contains(name)) {
|
|
extensionCatalog.add(cb);
|
|
if (enabled.contains(name)) {
|
|
active.add(cb);
|
|
}
|
|
} else {
|
|
active.add(cb);
|
|
}
|
|
}
|
|
return new ToolDisclosureSplit(active, extensionCatalog);
|
|
}
|
|
|
|
/**
|
|
* {@inheritDoc}
|
|
*
|
|
* <p>Protection set: only the {@link #ALWAYS_CORE} recovery/bridge tools.
|
|
* An explicit {@code disclosure_tier = core} remains a preference, but it
|
|
* cannot override the hard per-request schema ceiling. MCP tools default
|
|
* to EXTENSION (Move 5) so they only enter the CORE list when an operator
|
|
* explicitly sets {@code disclosure_tier = core} on the server — in that
|
|
* case they are still demotable, since MCP schemas are typically the
|
|
* heaviest part of the advertisement.
|
|
*/
|
|
@Override
|
|
public Set<String> computeAutoDemotions(AgentToolSet baseSet, Integer budgetTokens) {
|
|
if (legacyMode() || baseSet == null || budgetTokens == null
|
|
|| budgetTokens <= 0 || budgetTokens == Integer.MAX_VALUE) {
|
|
return Set.of();
|
|
}
|
|
List<ToolCallback> core = split(baseSet, Set.of()).activeCallbacks();
|
|
int coreTokens = TokenEstimator.estimateToolsTokens(core);
|
|
if (coreTokens <= budgetTokens) {
|
|
return Set.of();
|
|
}
|
|
Snapshot snap = snapshot();
|
|
List<ToolCallback> candidates = core.stream()
|
|
.filter(cb -> isDemotable(cb.getToolDefinition().name(), snap))
|
|
.sorted(demotionOrder())
|
|
.toList();
|
|
Set<String> demoted = new LinkedHashSet<>();
|
|
int remainingTokens = coreTokens;
|
|
for (ToolCallback cb : candidates) {
|
|
if (remainingTokens <= budgetTokens) {
|
|
break;
|
|
}
|
|
remainingTokens -= TokenEstimator.estimateToolsTokens(List.of(cb));
|
|
demoted.add(cb.getToolDefinition().name());
|
|
}
|
|
if (!demoted.isEmpty()) {
|
|
log.info("[ToolDisclosure] 工具 schema 估算 {} tokens 超出预算 {}——已将 {} 个最少使用的工具"
|
|
+ "降级到渐进目录(tool_call 可当轮调用): {}",
|
|
coreTokens, budgetTokens, demoted.size(), demoted);
|
|
}
|
|
return demoted;
|
|
}
|
|
|
|
private boolean isDemotable(String toolName, Snapshot snap) {
|
|
if (toolName == null || ALWAYS_CORE.contains(toolName)) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
/** Never-used tools demote first, then least recently used; name-tiebreak keeps builds deterministic. */
|
|
private Comparator<ToolCallback> demotionOrder() {
|
|
return Comparator
|
|
.<ToolCallback, Long>comparing(cb -> {
|
|
Long lastUsed = usageRecencyTracker == null
|
|
? null : usageRecencyTracker.lastUsedAt(cb.getToolDefinition().name());
|
|
return lastUsed == null ? Long.MIN_VALUE : lastUsed;
|
|
})
|
|
.thenComparing(cb -> cb.getToolDefinition().name());
|
|
}
|
|
|
|
@Override
|
|
public String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens) {
|
|
return renderExtensionCatalog(baseSet, maxInputTokens, Set.of());
|
|
}
|
|
|
|
@Override
|
|
public String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens,
|
|
Set<String> autoDemoted) {
|
|
if (legacyMode() || baseSet == null) {
|
|
return "";
|
|
}
|
|
List<ToolCallback> extension = split(baseSet, Set.of(), autoDemoted).extensionCatalog();
|
|
if (extension.isEmpty()) {
|
|
return "";
|
|
}
|
|
int limit = catalogEntryLimit(maxInputTokens);
|
|
Snapshot snap = snapshot();
|
|
|
|
StringBuilder sb = new StringBuilder();
|
|
sb.append("\n\n## Extension Tools\n");
|
|
sb.append("Full schemas are hidden until needed. If the exact name is known, call ");
|
|
sb.append("`tool_call(toolName=\"<name>\", arguments={...})` to execute it in this same round. ");
|
|
sb.append("Use `tool_search` to discover by capability and `tool_describe` only when arguments are unclear.\n\n");
|
|
sb.append("| Tool | Source | Description |\n");
|
|
sb.append("|------|--------|-------------|\n");
|
|
int shown = 0;
|
|
for (ToolCallback cb : extension) {
|
|
if (shown >= limit) break;
|
|
String name = cb.getToolDefinition().name();
|
|
sb.append("| `").append(name).append("` | ")
|
|
.append(sourceLabel(name, snap)).append(" | ");
|
|
String desc = cb.getToolDefinition().description();
|
|
if (desc != null && !desc.isBlank()) {
|
|
String d = desc.length() > 80 ? desc.substring(0, 80) + "..." : desc;
|
|
sb.append(d.replace("|", "\\|").replace("\n", " "));
|
|
}
|
|
sb.append(" |\n");
|
|
shown++;
|
|
}
|
|
if (extension.size() > shown) {
|
|
sb.append("\nShowing ").append(shown).append(" of ").append(extension.size())
|
|
.append(" extension tools.\n");
|
|
}
|
|
return sb.toString();
|
|
}
|
|
|
|
@Override
|
|
public void invalidate() {
|
|
this.snapshot = null;
|
|
}
|
|
|
|
private String sourceLabel(String toolName, Snapshot snap) {
|
|
Long serverId = snap.mcpToolToServerId.get(toolName);
|
|
if (serverId != null) {
|
|
String serverName = snap.serverNameById.get(serverId);
|
|
return serverName != null && !serverName.isBlank() ? "mcp:" + serverName : "mcp";
|
|
}
|
|
return "builtin";
|
|
}
|
|
|
|
private static int catalogEntryLimit(Integer maxInputTokens) {
|
|
if (maxInputTokens == null || maxInputTokens <= 0) return 20;
|
|
if (maxInputTokens < 8_000) return 12;
|
|
if (maxInputTokens < 32_000) return 25;
|
|
return 40;
|
|
}
|
|
|
|
private Snapshot snapshot() {
|
|
Snapshot snap = this.snapshot;
|
|
if (snap != null && (System.currentTimeMillis() - snap.builtAtMillis) < CACHE_TTL_MS) {
|
|
return snap;
|
|
}
|
|
Snapshot rebuilt = buildSnapshot();
|
|
this.snapshot = rebuilt;
|
|
return rebuilt;
|
|
}
|
|
|
|
private Snapshot buildSnapshot() {
|
|
// resolveTier() queries by the runtime function name (cb.getToolDefinition().name()),
|
|
// but mate_tool stores the Java class name (e.g. "ImageGenerateTool") and bean name
|
|
// (e.g. "imageGenerateTool"). Bridge both onto the function name(s) via a built-in
|
|
// alias index so a persisted tier actually reaches the runtime split. Do not call
|
|
// ToolRegistry.getEnabledToolSet() here: it synchronously enumerates MCP providers.
|
|
Map<String, DisclosureTier> builtinTierByName = new LinkedHashMap<>();
|
|
Map<String, Set<String>> builtinFunctionIndex = Map.of();
|
|
try {
|
|
builtinFunctionIndex = toolRegistry.enabledToolBeanFunctionNameIndex();
|
|
} catch (Exception e) {
|
|
log.warn("ToolDisclosureService: built-in tier name bridge unavailable: {}", e.getMessage());
|
|
}
|
|
try {
|
|
for (ToolEntity t : toolService.listTools()) {
|
|
if (t.getName() == null || t.getDisclosureTier() == null || t.getDisclosureTier().isBlank()) {
|
|
continue;
|
|
}
|
|
DisclosureTier tier = DisclosureTier.fromToken(t.getDisclosureTier());
|
|
// Key by the raw stored name too — harmless, and covers rows that already
|
|
// store a function name.
|
|
builtinTierByName.put(t.getName(), tier);
|
|
Set<String> aliases = new LinkedHashSet<>();
|
|
aliases.add(t.getName());
|
|
if (t.getBeanName() != null && !t.getBeanName().isBlank()) {
|
|
aliases.add(t.getBeanName());
|
|
}
|
|
for (String alias : aliases) {
|
|
Set<String> functionNames = builtinFunctionIndex.get(alias);
|
|
if (functionNames == null) {
|
|
continue;
|
|
}
|
|
for (String functionName : functionNames) {
|
|
builtinTierByName.put(functionName, tier);
|
|
}
|
|
}
|
|
}
|
|
} catch (Exception e) {
|
|
log.warn("ToolDisclosureService: failed to read mate_tool tiers, defaulting builtin tools to core: {}",
|
|
e.getMessage());
|
|
}
|
|
|
|
Map<String, Long> mcpToolToServerId = new LinkedHashMap<>();
|
|
try {
|
|
for (McpServerEntity server : mcpServerService.listEnabled()) {
|
|
if (server == null || server.getId() == null || server.getToolsCacheJson() == null
|
|
|| server.getToolsCacheJson().isBlank()) {
|
|
continue;
|
|
}
|
|
for (String toolName : cachedMcpToolNames(server)) {
|
|
mcpToolToServerId.put(toolName, server.getId());
|
|
}
|
|
}
|
|
} catch (Exception e) {
|
|
log.warn("ToolDisclosureService: failed to map MCP tools to servers: {}", e.getMessage());
|
|
}
|
|
|
|
Map<Long, DisclosureTier> serverTierById = new LinkedHashMap<>();
|
|
Map<Long, String> serverNameById = new LinkedHashMap<>();
|
|
try {
|
|
for (McpServerEntity s : mcpServerService.listAll()) {
|
|
// Move 5: only record an explicit tier. Servers with null
|
|
// disclosure_tier are intentionally absent from the map so
|
|
// resolveTierByName's getOrDefault(serverId, EXTENSION)
|
|
// applies the new on-demand default. Putting CORE here
|
|
// (via fromToken(null) → CORE) would override the default.
|
|
if (s.getDisclosureTier() != null && !s.getDisclosureTier().isBlank()) {
|
|
serverTierById.put(s.getId(), DisclosureTier.fromToken(s.getDisclosureTier()));
|
|
}
|
|
serverNameById.put(s.getId(), s.getName());
|
|
}
|
|
} catch (Exception e) {
|
|
log.warn("ToolDisclosureService: failed to read MCP server tiers, defaulting to extension: {}",
|
|
e.getMessage());
|
|
}
|
|
|
|
return new Snapshot(builtinTierByName, mcpToolToServerId, serverTierById, serverNameById,
|
|
System.currentTimeMillis());
|
|
}
|
|
|
|
private List<String> cachedMcpToolNames(McpServerEntity server) {
|
|
try {
|
|
JSONArray arr = JSONUtil.parseArray(server.getToolsCacheJson());
|
|
List<String> rawNames = new ArrayList<>(arr.size());
|
|
for (Object o : arr) {
|
|
if (!(o instanceof JSONObject jo)) {
|
|
continue;
|
|
}
|
|
String name = jo.getStr("name");
|
|
if (name != null && !name.isBlank()) {
|
|
rawNames.add(name);
|
|
}
|
|
}
|
|
if (rawNames.isEmpty()) {
|
|
return List.of();
|
|
}
|
|
return McpHashCollisionDetector.classify(server.getId(), rawNames).stream()
|
|
.filter(McpHashCollisionDetector.Decision::bindable)
|
|
.map(McpHashCollisionDetector.Decision::prefixedName)
|
|
.toList();
|
|
} catch (Exception e) {
|
|
log.debug("ToolDisclosureService: failed to parse MCP tools cache for server {}: {}",
|
|
server.getId(), e.getMessage());
|
|
return List.of();
|
|
}
|
|
}
|
|
|
|
private record Snapshot(Map<String, DisclosureTier> builtinTierByName,
|
|
Map<String, Long> mcpToolToServerId,
|
|
Map<Long, DisclosureTier> serverTierById,
|
|
Map<Long, String> serverNameById,
|
|
long builtAtMillis) {
|
|
}
|
|
}
|