diff --git a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java index 9437fa36..a3a5701a 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java @@ -52,6 +52,7 @@ import vip.mate.skill.runtime.SkillCatalogRenderer; import vip.mate.skill.service.SkillService; import vip.mate.system.service.SystemSettingService; import vip.mate.tool.ToolRegistry; +import vip.mate.tool.disclosure.ToolUsageRecencyTracker; import vip.mate.memory.spi.MemoryManager; import vip.mate.workspace.document.WorkspaceFileService; import vip.mate.tool.guard.service.ToolGuardService; @@ -103,6 +104,7 @@ public class AgentGraphBuilder { private final ModelProviderService modelProviderService; private final ModelContextWindowResolver contextWindowResolver; private final PrefixBudgetPlanner prefixBudgetPlanner; + private final ToolUsageRecencyTracker toolUsageRecencyTracker; private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService; private final ProviderRouter providerRouter; private final PlanningService planningService; @@ -428,10 +430,17 @@ public class AgentGraphBuilder { // in ReasoningNode; Plan-Execute keeps advertising every tool (it has no // action node to record enable_tool), so baking the catalog there would // describe an enable_tool flow that can never take effect. + // Auto-demotion is likewise ReAct-only: hiding a tool from Plan-Execute + // would remove it with no enable_tool path to recover it. boolean isPlanExecute = "plan_execute".equals(entity.getAgentType()); + Set autoDemotedTools = Set.of(); if (!isPlanExecute) { + if (prefixBudgetPlan.enabled()) { + autoDemotedTools = toolDisclosureService.computeAutoDemotions( + toolSet, prefixBudgetPlan.toolSchemaBudgetTokens()); + } String extensionCatalog = toolDisclosureService.renderExtensionCatalog( - toolSet, effectiveMaxInputTokens); + toolSet, effectiveMaxInputTokens, autoDemotedTools); if (extensionCatalog != null && !extensionCatalog.isBlank()) { enhancedPrompt = enhancedPrompt + extensionCatalog; } @@ -452,7 +461,7 @@ public class AgentGraphBuilder { entity.getName(), maxIter, toolSet.size(), protocol.getId()); } else { agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer, - prefixBudgetPlan); + prefixBudgetPlan, autoDemotedTools); // StateGraph 路径下工具调用由 ActionNode 控制，始终启用 toolCallingEnabled = true; log.info("Built StateGraph ReAct agent: {} (maxIterations={}, tools={}, protocol={})", @@ -539,17 +548,17 @@ public class AgentGraphBuilder { StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer) { - return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null); + return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null, Set.of()); } StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer, - PrefixBudgetPlan prefixBudgetPlan) { + PrefixBudgetPlan prefixBudgetPlan, Set autoDemotedTools) { ChatModel chatModel = buildRuntimeChatModel(runtimeModel); ChatClient chatClient = ChatClient.create(chatModel); String reasoningEffort = resolveReasoningEffortForModel(runtimeModel); CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort, - runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan); + runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan, autoDemotedTools); return new StateGraphReActAgent(chatClient, conversationService, compiledGraph, chatModel, conversationWindowManager, toolSet); } @@ -616,6 +625,7 @@ public class AgentGraphBuilder { // LLM mis-calls a skill name as a tool, the response tells it // the right invocation pattern instead of a dead-end error. executor.setSkillRuntimeService(skillRuntimeService); + executor.setUsageRecencyTracker(toolUsageRecencyTracker); // Optional: route child-agent denied-tool audit events through // the audit pipeline. Null when audit is not wired (legacy / test). if (auditEventService != null) { @@ -876,13 +886,13 @@ public class AgentGraphBuilder { String reasoningEffort, ModelConfigEntity primaryModelConfig, Long agentId, SkillCatalogRenderer skillCatalogRenderer) { return buildReActGraph(toolSet, chatModel, maxIterations, reasoningEffort, - primaryModelConfig, agentId, skillCatalogRenderer, null); + primaryModelConfig, agentId, skillCatalogRenderer, null, Set.of()); } CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, String reasoningEffort, ModelConfigEntity primaryModelConfig, Long agentId, SkillCatalogRenderer skillCatalogRenderer, - PrefixBudgetPlan prefixBudgetPlan) { + PrefixBudgetPlan prefixBudgetPlan, Set autoDemotedTools) { try { List fallbackChain = buildFallbackChain(primaryModelConfig, agentId); NodeStreamingChatHelper streamingHelper = new NodeStreamingChatHelper( @@ -905,6 +915,7 @@ public class AgentGraphBuilder { // LLM mis-calls a skill name as a tool, the response tells it // the right invocation pattern instead of a dead-end error. executor.setSkillRuntimeService(skillRuntimeService); + executor.setUsageRecencyTracker(toolUsageRecencyTracker); // Optional: route child-agent denied-tool audit events through // the audit pipeline. Null when audit is not wired (legacy / test). if (auditEventService != null) { @@ -920,6 +931,7 @@ public class AgentGraphBuilder { streamingHelper, conversationWindowManager, streamTracker, 0, wikiContextService, skillCatalogRenderer, toolDisclosureService, progressLedgerService); reasoningNode.setPrefixBudgetPlan(prefixBudgetPlan); + reasoningNode.setAutoDemotedTools(autoDemotedTools); ActionNode actionNode = new ActionNode(executor, streamTracker); ObservationProcessor observationProcessor = new ObservationProcessor(graphObservationProperties); ObservationNode observationNode = new ObservationNode(observationProcessor, streamTracker); diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlan.java b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlan.java index 638fa0d7..d0b71856 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlan.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlan.java @@ -19,7 +19,8 @@ public record PrefixBudgetPlan( int wikiTokens, int skillCatalogTokens, int extensionCatalogTokens, - int ledgerTokens) { + int ledgerTokens, + int toolSchemaBudgetTokens) { /** Window-size tier. Small windows tighten the injection ratio. */ public enum Profile { @@ -35,6 +36,6 @@ public record PrefixBudgetPlan( public static PrefixBudgetPlan unlimited(int effectiveMaxTokens) { return new PrefixBudgetPlan(false, effectiveMaxTokens, Profile.NORMAL, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, - Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE); + Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE, Integer.MAX_VALUE); } } diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlanner.java b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlanner.java index b393e651..1eda6a8f 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlanner.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlanner.java @@ -72,7 +72,8 @@ public class PrefixBudgetPlanner { (int) (injectionBudget * shares.getWiki() / sum), (int) (injectionBudget * shares.getSkill() / sum), (int) (injectionBudget * shares.getExtensionCatalog() / sum), - (int) (injectionBudget * shares.getLedger() / sum)); + (int) (injectionBudget * shares.getLedger() / sum), + (int) (effectiveMax * properties.getToolSchemaRatio())); if (profile != PrefixBudgetPlan.Profile.NORMAL) { log.info("[PrefixBudget] 窗口 {} tokens 进入 {} 档:注入预算 {} tokens" diff --git a/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java index 71363c40..765096c8 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/graph/executor/ToolExecutionExecutor.java @@ -7,6 +7,7 @@ import org.springframework.ai.chat.messages.ToolResponseMessage; import org.springframework.ai.chat.model.ToolContext; import org.springframework.ai.tool.ToolCallback; import vip.mate.tool.builtin.ToolExecutionContext; +import vip.mate.tool.disclosure.ToolUsageRecencyTracker; import vip.mate.agent.AgentToolSet; import vip.mate.agent.GraphEventPublisher; import vip.mate.agent.context.ChatOrigin; @@ -251,6 +252,13 @@ public class ToolExecutionExecutor { */ private vip.mate.skill.runtime.SkillRuntimeService skillRuntimeService; + /** Optional recency feed for budget-driven tool-disclosure demotion. */ + private ToolUsageRecencyTracker usageRecencyTracker; + + public void setUsageRecencyTracker(ToolUsageRecencyTracker tracker) { + this.usageRecencyTracker = tracker; + } + public void setSkillRuntimeService(vip.mate.skill.runtime.SkillRuntimeService s) { this.skillRuntimeService = s; } @@ -885,6 +893,12 @@ public class ToolExecutionExecutor { ToolExecutionContext.clear(); } + // Recency feed for budget-driven disclosure demotion: recently used + // tools keep their advertised schema, never-used ones demote first. + if (usageRecencyTracker != null) { + usageRecencyTracker.recordUse(toolName); + } + int rawLen = result != null ? result.length() : 0; // RFC-052: returnDirect tools bypass spill / truncation / LLM context. // Their full text goes to the user verbatim and is never persisted to diff --git a/mateclaw-server/src/main/java/vip/mate/agent/graph/node/ReasoningNode.java b/mateclaw-server/src/main/java/vip/mate/agent/graph/node/ReasoningNode.java index 6e0ed0ab..2cd48158 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/graph/node/ReasoningNode.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/graph/node/ReasoningNode.java @@ -363,6 +363,19 @@ public class ReasoningNode implements NodeAction { this.prefixBudgetPlan = prefixBudgetPlan; } + /** + * Core-tier tools auto-demoted to the extension catalog because the + * advertised schemas exceeded the window's tool-schema budget. Decided + * once at agent-build time (kept stable for prompt caching); the baked + * extension catalog lists them so {@code enable_tool} can surface any of + * them back. + */ + private Set autoDemotedTools = Set.of(); + + public void setAutoDemotedTools(Set autoDemotedTools) { + this.autoDemotedTools = autoDemotedTools == null ? Set.of() : autoDemotedTools; + } + public ReasoningNode(ChatModel chatModel, AgentToolSet toolSet, String reasoningEffort, NodeStreamingChatHelper streamingHelper, ConversationWindowManager conversationWindowManager, @@ -728,7 +741,8 @@ public class ReasoningNode implements NodeAction { // so an enable_tool call earlier in this loop takes effect immediately. // Falls back to the full tool set when no disclosure service is wired. List activeCallbacks = (toolDisclosureService != null && toolSet != null) - ? toolDisclosureService.split(toolSet, accessor.enabledExtensionTools()).activeCallbacks() + ? toolDisclosureService.split(toolSet, accessor.enabledExtensionTools(), autoDemotedTools) + .activeCallbacks() : toolCallbacks; ChatOptions options = buildChatOptions(effectiveReasoning, activeCallbacks); diff --git a/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java b/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java index 60c84274..0bcd5dc2 100644 --- a/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java +++ b/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java @@ -47,6 +47,14 @@ public class PrefixBudgetProperties { */ private double compactTriggerRatioOverride = 0.85; + /** + * Fraction of the effective window the advertised tool schemas may + * occupy. When the core tool set estimates above this, the least + * recently used demotable tools are auto-moved to the extension catalog + * (recoverable via {@code enable_tool}) until the set fits. + */ + private double toolSchemaRatio = 0.25; + /** Relative shares of the injection budget. Normalized at plan time. */ private Shares shares = new Shares(); diff --git a/mateclaw-server/src/main/java/vip/mate/tool/disclosure/DefaultToolDisclosureService.java b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/DefaultToolDisclosureService.java index 81120b40..b10345cb 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/disclosure/DefaultToolDisclosureService.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/DefaultToolDisclosureService.java @@ -6,6 +6,7 @@ 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.service.McpServerService; @@ -15,6 +16,7 @@ import vip.mate.tool.service.AvailableToolService; 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; @@ -53,6 +55,7 @@ public class DefaultToolDisclosureService implements ToolDisclosureService { private final McpServerService mcpServerService; private final AvailableToolService availableToolService; private final ToolRegistry toolRegistry; + private final ToolUsageRecencyTracker usageRecencyTracker; @Value("${mateclaw.tools.disclosure.mode:progressive}") private String disclosureMode; @@ -97,17 +100,25 @@ public class DefaultToolDisclosureService implements ToolDisclosureService { @Override public ToolDisclosureSplit split(AgentToolSet baseSet, Set enabledExtensions) { + return split(baseSet, enabledExtensions, Set.of()); + } + + @Override + public ToolDisclosureSplit split(AgentToolSet baseSet, Set enabledExtensions, + Set autoDemoted) { List all = baseSet == null ? List.of() : baseSet.callbacks(); if (legacyMode()) { return new ToolDisclosureSplit(all, List.of()); } Set enabled = enabledExtensions == null ? Set.of() : enabledExtensions; + Set demoted = autoDemoted == null ? Set.of() : autoDemoted; List active = new ArrayList<>(all.size()); List extensionCatalog = new ArrayList<>(); for (ToolCallback cb : all) { - if (resolveTier(cb) == DisclosureTier.EXTENSION) { + String name = cb.getToolDefinition().name(); + if (resolveTier(cb) == DisclosureTier.EXTENSION || demoted.contains(name)) { extensionCatalog.add(cb); - if (enabled.contains(cb.getToolDefinition().name())) { + if (enabled.contains(name)) { active.add(cb); } } else { @@ -117,12 +128,79 @@ public class DefaultToolDisclosureService implements ToolDisclosureService { return new ToolDisclosureSplit(active, extensionCatalog); } + /** + * {@inheritDoc} + * + *

Protection set: {@link #ALWAYS_CORE} meta-tools and builtin tools + * with an explicit {@code disclosure_tier = core} row. MCP tools remain + * demotable — the server-level tier cannot distinguish an explicit core + * choice from the default, and MCP schemas are typically the heaviest + * part of the advertisement. + */ + @Override + public Set computeAutoDemotions(AgentToolSet baseSet, Integer budgetTokens) { + if (legacyMode() || baseSet == null || budgetTokens == null + || budgetTokens <= 0 || budgetTokens == Integer.MAX_VALUE) { + return Set.of(); + } + List core = split(baseSet, Set.of()).activeCallbacks(); + int coreTokens = TokenEstimator.estimateToolsTokens(core); + if (coreTokens <= budgetTokens) { + return Set.of(); + } + Snapshot snap = snapshot(); + List candidates = core.stream() + .filter(cb -> isDemotable(cb.getToolDefinition().name(), snap)) + .sorted(demotionOrder()) + .toList(); + Set 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 超出预算 {}——已将 {} 个最少使用的工具" + + "降级到扩展目录(enable_tool 可找回): {}", + coreTokens, budgetTokens, demoted.size(), demoted); + } + return demoted; + } + + private boolean isDemotable(String toolName, Snapshot snap) { + if (toolName == null || ALWAYS_CORE.contains(toolName)) { + return false; + } + // An explicit core row is an operator decision — never override it. + return snap.builtinTierByName.get(toolName) != DisclosureTier.CORE; + } + + /** Never-used tools demote first, then least recently used; name-tiebreak keeps builds deterministic. */ + private Comparator demotionOrder() { + return Comparator + .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 autoDemoted) { if (legacyMode() || baseSet == null) { return ""; } - List extension = split(baseSet, Set.of()).extensionCatalog(); + List extension = split(baseSet, Set.of(), autoDemoted).extensionCatalog(); if (extension.isEmpty()) { return ""; } diff --git a/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolDisclosureService.java b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolDisclosureService.java index 9ad06d83..c1b16b84 100644 --- a/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolDisclosureService.java +++ b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolDisclosureService.java @@ -31,6 +31,29 @@ public interface ToolDisclosureService { */ ToolDisclosureSplit split(AgentToolSet baseSet, Set enabledExtensions); + /** + * Budget-aware variant: tools in {@code autoDemoted} are treated as + * extension tier for this split even when their resolved tier is core. + * The demotion set is decided once per agent build (see + * {@link #computeAutoDemotions}) so the runtime split, the baked catalog + * and the prompt-cache prefix stay consistent with each other. + */ + default ToolDisclosureSplit split(AgentToolSet baseSet, Set enabledExtensions, + Set autoDemoted) { + return split(baseSet, enabledExtensions); + } + + /** + * Decide which core-tier tools to auto-demote so the advertised tool + * schemas fit {@code budgetTokens} (estimated). Ranking: never-used tools + * first, then least recently used; meta-tools and explicitly configured + * core tools are never demoted. Empty when the set already fits, when + * {@code budgetTokens} is null, or in legacy disclosure mode. + */ + default Set computeAutoDemotions(AgentToolSet baseSet, Integer budgetTokens) { + return Set.of(); + } + /** * Render the {@code ## Extension Tools} system-prompt segment for the * agent's extension tools, or an empty string when there are none / when @@ -38,6 +61,15 @@ public interface ToolDisclosureService { */ String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens); + /** + * Budget-aware variant: auto-demoted tools are listed in the catalog too, + * so the model can discover and {@code enable_tool} them back. + */ + default String renderExtensionCatalog(AgentToolSet baseSet, Integer maxInputTokens, + Set autoDemoted) { + return renderExtensionCatalog(baseSet, maxInputTokens); + } + /** Drop the cached tier snapshot so the next resolve re-reads the DB. */ void invalidate(); diff --git a/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolUsageRecencyTracker.java b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolUsageRecencyTracker.java new file mode 100644 index 00000000..f01f6a6a --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/tool/disclosure/ToolUsageRecencyTracker.java @@ -0,0 +1,33 @@ +package vip.mate.tool.disclosure; + +import org.springframework.stereotype.Component; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * In-memory recency signal for tool usage. Feeds the budget-driven + * auto-demotion ranking: never-used tools demote first, then the least + * recently used ones. + * + *

Deliberately process-local and unpersisted — this is an advisory + * ranking, not an audit trail. A restart resets everything to "never used", + * which merely makes the first demotion pass alphabetical. + */ +@Component +public class ToolUsageRecencyTracker { + + private final Map lastUsedAtMs = new ConcurrentHashMap<>(); + + /** Record a successful execution of {@code toolName}. */ + public void recordUse(String toolName) { + if (toolName != null && !toolName.isBlank()) { + lastUsedAtMs.put(toolName, System.currentTimeMillis()); + } + } + + /** @return epoch millis of the last recorded use, or null when never used. */ + public Long lastUsedAt(String toolName) { + return toolName == null ? null : lastUsedAtMs.get(toolName); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/agent/graph/node/ReasoningNodePtlPromptTest.java b/mateclaw-server/src/test/java/vip/mate/agent/graph/node/ReasoningNodePtlPromptTest.java index 3c0d4851..1d577ce3 100644 --- a/mateclaw-server/src/test/java/vip/mate/agent/graph/node/ReasoningNodePtlPromptTest.java +++ b/mateclaw-server/src/test/java/vip/mate/agent/graph/node/ReasoningNodePtlPromptTest.java @@ -12,6 +12,7 @@ import java.util.List; import static org.assertj.core.api.Assertions.assertThat; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -38,7 +39,7 @@ class ReasoningNodePtlPromptTest { @Test void prefixIncludesSystemRuntimeAndWikiSegments() { WikiContextService wikiContextService = mock(WikiContextService.class); - when(wikiContextService.buildRelevantContext(eq(42L), anyString())) + when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull())) .thenReturn(WIKI_RELEVANT_TEXT); ReasoningNode node = newNode(wikiContextService); @@ -67,7 +68,7 @@ class ReasoningNodePtlPromptTest { // that two independent calls with the same inputs produce // structurally identical output. WikiContextService wikiContextService = mock(WikiContextService.class); - when(wikiContextService.buildRelevantContext(eq(42L), anyString())) + when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull())) .thenReturn(WIKI_RELEVANT_TEXT); ReasoningNode node = newNode(wikiContextService); @@ -125,7 +126,7 @@ class ReasoningNodePtlPromptTest { @Test void blankWikiResultSkipsWikiSegment() { WikiContextService wikiContextService = mock(WikiContextService.class); - when(wikiContextService.buildRelevantContext(eq(42L), anyString())) + when(wikiContextService.buildRelevantContext(eq(42L), anyString(), isNull())) .thenReturn(" "); // blank → drop the layer ReasoningNode node = newNode(wikiContextService); diff --git a/mateclaw-server/src/test/java/vip/mate/tool/disclosure/ToolDisclosureServiceTest.java b/mateclaw-server/src/test/java/vip/mate/tool/disclosure/ToolDisclosureServiceTest.java index 15ba8aad..bd343fb6 100644 --- a/mateclaw-server/src/test/java/vip/mate/tool/disclosure/ToolDisclosureServiceTest.java +++ b/mateclaw-server/src/test/java/vip/mate/tool/disclosure/ToolDisclosureServiceTest.java @@ -7,6 +7,7 @@ import org.springframework.ai.tool.ToolCallback; import org.springframework.ai.tool.annotation.Tool; import org.springframework.test.util.ReflectionTestUtils; 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.service.McpServerService; @@ -88,7 +89,7 @@ class ToolDisclosureServiceTest { lenient().when(ms.listAll()).thenReturn(servers); lenient().when(as.listAvailable()).thenReturn(available); lenient().when(tr.getEnabledToolSet()).thenReturn(globalSet()); - return new DefaultToolDisclosureService(ts, ms, as, tr); + return new DefaultToolDisclosureService(ts, ms, as, tr, new ToolUsageRecencyTracker()); } @Test @@ -202,4 +203,105 @@ class ToolDisclosureServiceTest { private static List names(List cbs) { return cbs.stream().map(c -> c.getToolDefinition().name()).toList(); } + + // ==================== budget-driven auto-demotion ==================== + + /** Three plain core tools for demotion-ranking tests. */ + static class ManyCoreTools { + @Tool(description = "core tool a") + public String tool_a() { return ""; } + + @Tool(description = "core tool b") + public String tool_b() { return ""; } + + @Tool(description = "core tool c") + public String tool_c() { return ""; } + } + + private static AgentToolSet manyCoreSet() { + return AgentToolSet.fromCallbacks(List.of(new ManyCoreTools()), + List.of(ToolCallbacks.from(new ManyCoreTools()))); + } + + @Test + @DisplayName("no demotion when the core schemas fit the budget, or when budget is absent") + void noDemotionWhenBudgetFits() { + var svc = service(List.of(), List.of(), List.of()); + AgentToolSet set = manyCoreSet(); + assertTrue(svc.computeAutoDemotions(set, Integer.MAX_VALUE).isEmpty()); + assertTrue(svc.computeAutoDemotions(set, null).isEmpty()); + assertTrue(svc.computeAutoDemotions(set, 1_000_000).isEmpty()); + } + + @Test + @DisplayName("tiny budget demotes every demotable tool, alphabetical when nothing was ever used") + void tinyBudgetDemotesAll() { + var svc = service(List.of(), List.of(), List.of()); + var demoted = svc.computeAutoDemotions(manyCoreSet(), 1); + assertEquals(Set.of("tool_a", "tool_b", "tool_c"), demoted); + } + + @Test + @DisplayName("budget one tool short demotes exactly the first never-used candidate") + void partialDemotionTakesFirstCandidate() { + var svc = service(List.of(), List.of(), List.of()); + AgentToolSet set = manyCoreSet(); + int coreTokens = TokenEstimator.estimateToolsTokens(svc.split(set, Set.of()).activeCallbacks()); + var demoted = svc.computeAutoDemotions(set, coreTokens - 1); + assertEquals(Set.of("tool_a"), demoted); + } + + @Test + @DisplayName("recently used tools demote last") + void recencyProtectsRecentlyUsed() { + ToolUsageRecencyTracker tracker = new ToolUsageRecencyTracker(); + tracker.recordUse("tool_a"); + ToolService ts = mock(ToolService.class); + McpServerService ms = mock(McpServerService.class); + AvailableToolService as = mock(AvailableToolService.class); + ToolRegistry tr = mock(ToolRegistry.class); + lenient().when(ts.listTools()).thenReturn(List.of()); + lenient().when(ms.listAll()).thenReturn(List.of()); + lenient().when(as.listAvailable()).thenReturn(List.of()); + lenient().when(tr.getEnabledToolSet()).thenReturn(globalSet()); + var svc = new DefaultToolDisclosureService(ts, ms, as, tr, tracker); + + AgentToolSet set = manyCoreSet(); + int coreTokens = TokenEstimator.estimateToolsTokens(svc.split(set, Set.of()).activeCallbacks()); + // One tool over budget: the never-used tool_b (alphabetically first + // among never-used) demotes, the recently used tool_a survives. + var demoted = svc.computeAutoDemotions(set, coreTokens - 1); + assertEquals(Set.of("tool_b"), demoted); + } + + @Test + @DisplayName("explicit core DB row and meta-tools are never demoted") + void explicitCoreProtected() { + var svc = service(List.of(toolRow("tool_a", "builtin", "core")), List.of(), List.of()); + var demoted = svc.computeAutoDemotions(manyCoreSet(), 1); + assertEquals(Set.of("tool_b", "tool_c"), demoted); + } + + @Test + @DisplayName("auto-demoted tools behave as extension in split and can be enabled back") + void splitHonorsAutoDemotions() { + var svc = service(List.of(), List.of(), List.of()); + AgentToolSet set = manyCoreSet(); + + var split = svc.split(set, Set.of(), Set.of("tool_b")); + assertEquals(List.of("tool_a", "tool_c"), names(split.activeCallbacks())); + assertEquals(List.of("tool_b"), names(split.extensionCatalog())); + + var enabledBack = svc.split(set, Set.of("tool_b"), Set.of("tool_b")); + assertTrue(names(enabledBack.activeCallbacks()).contains("tool_b")); + } + + @Test + @DisplayName("catalog rendering lists auto-demoted tools for discoverability") + void catalogListsAutoDemoted() { + var svc = service(List.of(), List.of(), List.of()); + String catalog = svc.renderExtensionCatalog(manyCoreSet(), 8192, Set.of("tool_b")); + assertTrue(catalog.contains("tool_b")); + assertFalse(catalog.contains("| `tool_a`"), "non-demoted core tools stay out of the catalog"); + } }