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 44183381..9437fa36 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java @@ -38,6 +38,9 @@ import vip.mate.llm.chatmodel.ReasoningEffortResolver; import vip.mate.llm.model.ModelConfigEntity; import vip.mate.llm.model.ModelFamily; import vip.mate.llm.model.ModelProtocol; +import vip.mate.agent.context.PrefixBudgetPlan; +import vip.mate.agent.context.PrefixBudgetPlanner; +import vip.mate.agent.context.TokenEstimator; import vip.mate.llm.model.ModelProviderEntity; import vip.mate.llm.probe.ModelContextWindowResolver; import vip.mate.llm.routing.ProviderModelRef; @@ -99,6 +102,7 @@ public class AgentGraphBuilder { private final ModelConfigService modelConfigService; private final ModelProviderService modelProviderService; private final ModelContextWindowResolver contextWindowResolver; + private final PrefixBudgetPlanner prefixBudgetPlanner; private final vip.mate.llm.service.ModelCapabilityService modelCapabilityService; private final ProviderRouter providerRouter; private final PlanningService planningService; @@ -396,7 +400,22 @@ public class AgentGraphBuilder { } } - String enhancedPrompt = buildEnhancedPrompt(entity, builtinSearchEnabled); + // Prefix injection budget: optional blocks (memory / wiki / skill + // catalog / extension catalog / ledger) share a token budget scaled + // to the model's effective window. The agent's own prompt and the + // tool schemas are never truncated — they are subtracted from the + // budget so the optional blocks absorb the squeeze. + int basePromptTokens = TokenEstimator.estimateTokens(entity.getSystemPrompt()); + int toolSchemaTokens = TokenEstimator.estimateToolsTokens(toolSet.callbacks()); + PrefixBudgetPlan prefixBudgetPlan = prefixBudgetPlanner.plan( + effectiveMaxInputTokens, basePromptTokens, toolSchemaTokens); + if (basePromptTokens > prefixBudgetPlan.effectiveMaxTokens() / 2) { + log.warn("Agent {} 的身份 prompt 约 {} tokens,已超过模型有效窗口 {} 的一半——" + + "系统不会截断用户自写的身份 prompt,请自行精简,否则小上下文模型可能无法响应", + entity.getId(), basePromptTokens, prefixBudgetPlan.effectiveMaxTokens()); + } + + String enhancedPrompt = buildEnhancedPrompt(entity, builtinSearchEnabled, prefixBudgetPlan.memoryTokens()); // Runtime skill-catalog renderer — captures this agent's bound skills, // effective tool allowlist, model window and workspace; invoked each @@ -432,7 +451,8 @@ public class AgentGraphBuilder { log.info("Built StateGraph Plan-Execute agent: {} (maxIterations={}, tools={}, protocol={})", entity.getName(), maxIter, toolSet.size(), protocol.getId()); } else { - agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer); + agent = buildReActAgent(toolSet, runtimeModel, maxIter, entity.getId(), skillCatalogRenderer, + prefixBudgetPlan); // StateGraph 路径下工具调用由 ActionNode 控制，始终启用 toolCallingEnabled = true; log.info("Built StateGraph ReAct agent: {} (maxIterations={}, tools={}, protocol={})", @@ -519,11 +539,17 @@ public class AgentGraphBuilder { StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer) { + return buildReActAgent(toolSet, runtimeModel, maxIter, agentId, skillCatalogRenderer, null); + } + + StateGraphReActAgent buildReActAgent(AgentToolSet toolSet, ModelConfigEntity runtimeModel, + int maxIter, Long agentId, SkillCatalogRenderer skillCatalogRenderer, + PrefixBudgetPlan prefixBudgetPlan) { ChatModel chatModel = buildRuntimeChatModel(runtimeModel); ChatClient chatClient = ChatClient.create(chatModel); String reasoningEffort = resolveReasoningEffortForModel(runtimeModel); CompiledGraph compiledGraph = buildReActGraph(toolSet, chatModel, maxIter, reasoningEffort, - runtimeModel, agentId, skillCatalogRenderer); + runtimeModel, agentId, skillCatalogRenderer, prefixBudgetPlan); return new StateGraphReActAgent(chatClient, conversationService, compiledGraph, chatModel, conversationWindowManager, toolSet); } @@ -849,6 +875,14 @@ public class AgentGraphBuilder { CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, String reasoningEffort, ModelConfigEntity primaryModelConfig, Long agentId, SkillCatalogRenderer skillCatalogRenderer) { + return buildReActGraph(toolSet, chatModel, maxIterations, reasoningEffort, + primaryModelConfig, agentId, skillCatalogRenderer, null); + } + + CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, + String reasoningEffort, ModelConfigEntity primaryModelConfig, + Long agentId, SkillCatalogRenderer skillCatalogRenderer, + PrefixBudgetPlan prefixBudgetPlan) { try { List fallbackChain = buildFallbackChain(primaryModelConfig, agentId); NodeStreamingChatHelper streamingHelper = new NodeStreamingChatHelper( @@ -885,6 +919,7 @@ public class AgentGraphBuilder { supportsReasoningEffort, streamingHelper, conversationWindowManager, streamTracker, 0, wikiContextService, skillCatalogRenderer, toolDisclosureService, progressLedgerService); + reasoningNode.setPrefixBudgetPlan(prefixBudgetPlan); ActionNode actionNode = new ActionNode(executor, streamTracker); ObservationProcessor observationProcessor = new ObservationProcessor(graphObservationProperties); ObservationNode observationNode = new ObservationNode(observationProcessor, streamTracker); @@ -1473,6 +1508,10 @@ public class AgentGraphBuilder { """; private String buildEnhancedPrompt(AgentEntity entity, boolean builtinSearchEnabled) { + return buildEnhancedPrompt(entity, builtinSearchEnabled, Integer.MAX_VALUE); + } + + private String buildEnhancedPrompt(AgentEntity entity, boolean builtinSearchEnabled, int memoryBudgetTokens) { // The agent's own systemPrompt encodes its identity (role / goal / // backstory). The memory block from workspace files (AGENTS.md, SOUL.md, // PROFILE.md, MEMORY.md, ...) augments that identity with durable @@ -1481,7 +1520,7 @@ public class AgentGraphBuilder { // dropped the identity prompt, so editor-side identity changes never // reached runtime if the agent had any workspace files. String identityPrompt = entity.getSystemPrompt() != null ? entity.getSystemPrompt().trim() : ""; - String memoryPrompt = memoryManager.buildSystemPromptBlock(entity.getId()); + String memoryPrompt = memoryManager.buildSystemPromptBlock(entity.getId(), memoryBudgetTokens); StringBuilder basePromptBuilder = new StringBuilder(); if (!identityPrompt.isEmpty()) { basePromptBuilder.append(identityPrompt); 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 new file mode 100644 index 00000000..638fa0d7 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlan.java @@ -0,0 +1,40 @@ +package vip.mate.agent.context; + +/** + * Per-agent-build token budget for the prompt prefix's optional injection + * blocks. Produced once by {@link PrefixBudgetPlanner} when the agent graph + * is assembled (the inputs — effective window, base prompt, tool schemas — + * are all stable per build) and handed to each injection site. + * + *

{@code enabled == false} means budgeting is switched off: every budget + * field holds {@link Integer#MAX_VALUE} and consumers keep their existing + * absolute caps untouched. + */ +public record PrefixBudgetPlan( + boolean enabled, + int effectiveMaxTokens, + Profile profile, + int injectionBudgetTokens, + int memoryTokens, + int wikiTokens, + int skillCatalogTokens, + int extensionCatalogTokens, + int ledgerTokens) { + + /** Window-size tier. Small windows tighten the injection ratio. */ + public enum Profile { + /** Regular window — budget shares rarely bind (absolute caps are smaller). */ + NORMAL, + /** Window below the compact threshold — tightened injection ratio. */ + COMPACT, + /** Window below the minimal threshold — injection cut to the bone. */ + MINIMAL + } + + /** Budgeting disabled — unlimited budgets, previous behavior. */ + 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); + } +} 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 new file mode 100644 index 00000000..b393e651 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/PrefixBudgetPlanner.java @@ -0,0 +1,105 @@ +package vip.mate.agent.context; + +import lombok.RequiredArgsConstructor; +import lombok.extern.slf4j.Slf4j; +import org.springframework.boot.context.properties.EnableConfigurationProperties; +import org.springframework.stereotype.Component; +import vip.mate.config.ConversationWindowProperties; +import vip.mate.config.PrefixBudgetProperties; + +/** + * Computes the {@link PrefixBudgetPlan} for one agent build: how many tokens + * each optional prefix injection block (memory / wiki / skill catalog / + * extension catalog / progress ledger) may spend, scaled to the model's + * effective context window. + * + * 

+ * injectionBudget = max(0, effectiveMax × ratio(profile)
+ *                          − basePromptTokens − toolSchemaTokens)
+ * block budget    = injectionBudget × normalizedShare(block)
+ *
+ * + * The agent's own prompt and the tool schemas are never truncated here — + * they are subtracted from the injection budget so the optional blocks + * absorb the squeeze. On large windows the shares far exceed each block's + * absolute cap, so behavior is byte-identical to the pre-budget code. + */ +@Slf4j +@Component +@RequiredArgsConstructor +@EnableConfigurationProperties(PrefixBudgetProperties.class) +public class PrefixBudgetPlanner { + + private final PrefixBudgetProperties properties; + private final ConversationWindowProperties windowProperties; + + /** + * @param effectiveMaxInputTokens the model's effective window (explicit + * config or probed); null/0 falls back to + * the global default + * @param basePromptTokens estimated tokens of the agent's own + * identity prompt (before memory/guidance) + * @param toolSchemaTokens estimated tokens of the advertised tool + * schemas + */ + public PrefixBudgetPlan plan(Integer effectiveMaxInputTokens, int basePromptTokens, int toolSchemaTokens) { + int effectiveMax = (effectiveMaxInputTokens != null && effectiveMaxInputTokens > 0) + ? effectiveMaxInputTokens : windowProperties.getDefaultMaxInputTokens(); + if (!properties.isEnabled()) { + return PrefixBudgetPlan.unlimited(effectiveMax); + } + + PrefixBudgetPlan.Profile profile = profileFor(effectiveMax); + double ratio = switch (profile) { + case NORMAL -> properties.getInjectionRatio(); + case COMPACT -> properties.getCompactInjectionRatio(); + case MINIMAL -> properties.getMinimalInjectionRatio(); + }; + + int injectionBudget = Math.max(0, + (int) (effectiveMax * ratio) - Math.max(0, basePromptTokens) - Math.max(0, toolSchemaTokens)); + + PrefixBudgetProperties.Shares shares = properties.getShares(); + double sum = shares.getMemory() + shares.getWiki() + shares.getSkill() + + shares.getExtensionCatalog() + shares.getLedger(); + if (sum <= 0) { + sum = 1.0; + } + + PrefixBudgetPlan plan = new PrefixBudgetPlan( + true, effectiveMax, profile, injectionBudget, + (int) (injectionBudget * shares.getMemory() / sum), + (int) (injectionBudget * shares.getWiki() / sum), + (int) (injectionBudget * shares.getSkill() / sum), + (int) (injectionBudget * shares.getExtensionCatalog() / sum), + (int) (injectionBudget * shares.getLedger() / sum)); + + if (profile != PrefixBudgetPlan.Profile.NORMAL) { + log.info("[PrefixBudget] 窗口 {} tokens 进入 {} 档:注入预算 {} tokens" + + "(memory={}, wiki={}, skill={}, extCatalog={}, ledger={})", + effectiveMax, profile, injectionBudget, + plan.memoryTokens(), plan.wikiTokens(), plan.skillCatalogTokens(), + plan.extensionCatalogTokens(), plan.ledgerTokens()); + } + return plan; + } + + /** Compaction trigger ratio for this window size (small windows fill up before compacting). */ + public double compactTriggerRatioFor(int effectiveMaxTokens, double defaultRatio) { + if (!properties.isEnabled()) { + return defaultRatio; + } + return profileFor(effectiveMaxTokens) == PrefixBudgetPlan.Profile.NORMAL + ? defaultRatio : properties.getCompactTriggerRatioOverride(); + } + + private PrefixBudgetPlan.Profile profileFor(int effectiveMax) { + if (effectiveMax < properties.getMinimalThresholdTokens()) { + return PrefixBudgetPlan.Profile.MINIMAL; + } + if (effectiveMax < properties.getCompactThresholdTokens()) { + return PrefixBudgetPlan.Profile.COMPACT; + } + return PrefixBudgetPlan.Profile.NORMAL; + } +} 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 0ec0da7d..6e0ed0ab 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 @@ -22,6 +22,7 @@ import vip.mate.llm.chatmodel.ThinkingLevelHolder; import vip.mate.agent.graph.NodeStreamingChatHelper; import vip.mate.agent.context.ConversationWindowManager; import vip.mate.agent.context.LoopBudgetConfig; +import vip.mate.agent.context.PrefixBudgetPlan; import vip.mate.agent.context.LoopMessageBudgeter; import vip.mate.agent.context.RuntimeContextInjector; import vip.mate.agent.context.TokenEstimator; @@ -350,6 +351,18 @@ public class ReasoningNode implements NodeAction { */ private final vip.mate.agent.progress.ProgressLedgerService progressLedgerService; + /** + * Token budget for the optional prefix injection blocks, computed at + * agent-build time against the model's effective context window. Null + * when the graph was assembled without budgeting (tests, legacy paths) — + * all injection sites then keep their previous absolute-cap behavior. + */ + private PrefixBudgetPlan prefixBudgetPlan; + + public void setPrefixBudgetPlan(PrefixBudgetPlan prefixBudgetPlan) { + this.prefixBudgetPlan = prefixBudgetPlan; + } + public ReasoningNode(ChatModel chatModel, AgentToolSet toolSet, String reasoningEffort, NodeStreamingChatHelper streamingHelper, ConversationWindowManager conversationWindowManager, @@ -470,6 +483,12 @@ public class ReasoningNode implements NodeAction { * otherwise the documented fallback. */ private int loopContextWindowTokens() { + // Per-model effective window (explicit config or probed) beats the + // global default — the loop budgeter is otherwise blind to small + // local models and never trims for them. + if (prefixBudgetPlan != null && prefixBudgetPlan.effectiveMaxTokens() > 0) { + return prefixBudgetPlan.effectiveMaxTokens(); + } if (conversationWindowManager != null) { int v = conversationWindowManager.getDefaultMaxInputTokens(); if (v > 0) return v; @@ -1119,7 +1138,9 @@ public class ReasoningNode implements NodeAction { if (!projectRecalled && wikiContextService != null && agentIdStr != null && !agentIdStr.isEmpty()) { try { Long parsedAgentId = Long.parseLong(agentIdStr); - String wikiRelevant = wikiContextService.buildRelevantContext(parsedAgentId, userMsg); + Integer wikiBudgetTokens = (prefixBudgetPlan != null && prefixBudgetPlan.enabled()) + ? prefixBudgetPlan.wikiTokens() : null; + String wikiRelevant = wikiContextService.buildRelevantContext(parsedAgentId, userMsg, wikiBudgetTokens); if (wikiRelevant != null && !wikiRelevant.isBlank()) { prefix.add(new UserMessage(wikiRelevant)); } diff --git a/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java b/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java new file mode 100644 index 00000000..60c84274 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/config/PrefixBudgetProperties.java @@ -0,0 +1,61 @@ +package vip.mate.config; + +import lombok.Data; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * Budget configuration for the prompt prefix's optional injection blocks + * (memory, wiki relevance, skill catalog, extension-tool catalog, progress + * ledger). + * + *

Each block already has its own absolute cap (chars), but those caps are + * mutually blind and sized for large cloud models — stacked together they + * easily exceed a local 8k/16k window on the very first request, when there + * is no history to compact. This budget scales every block against the + * model's effective context window instead: the enforced limit is always + * {@code min(block's own cap, its share of the injection budget)}, so large + * windows behave exactly as before while small windows shrink each block + * proportionally. + */ +@Data +@ConfigurationProperties(prefix = "mateclaw.context.prefix-budget") +public class PrefixBudgetProperties { + + /** Kill switch — false restores the previous per-block-absolute-cap behavior. */ + private boolean enabled = true; + + /** Fraction of the effective window granted to prefix injection blocks (normal profile). */ + private double injectionRatio = 0.35; + + /** Injection ratio when the window is below {@link #compactThresholdTokens}. */ + private double compactInjectionRatio = 0.25; + + /** Injection ratio when the window is below {@link #minimalThresholdTokens}. */ + private double minimalInjectionRatio = 0.15; + + /** Windows below this enter the compact profile. */ + private int compactThresholdTokens = 32768; + + /** Windows below this enter the minimal profile. */ + private int minimalThresholdTokens = 8192; + + /** + * Compaction trigger ratio used for compact / minimal profiles instead of + * the global {@code compactTriggerRatio}. Small windows should be used up + * before summarizing — compacting at 75% of an 8k window wastes what + * little room there is. + */ + private double compactTriggerRatioOverride = 0.85; + + /** Relative shares of the injection budget. Normalized at plan time. */ + private Shares shares = new Shares(); + + @Data + public static class Shares { + private double memory = 0.35; + private double wiki = 0.30; + private double skill = 0.20; + private double extensionCatalog = 0.10; + private double ledger = 0.05; + } +} diff --git a/mateclaw-server/src/main/java/vip/mate/memory/spi/MemoryManager.java b/mateclaw-server/src/main/java/vip/mate/memory/spi/MemoryManager.java index e0ba36ff..bbbcf0b7 100644 --- a/mateclaw-server/src/main/java/vip/mate/memory/spi/MemoryManager.java +++ b/mateclaw-server/src/main/java/vip/mate/memory/spi/MemoryManager.java @@ -3,6 +3,7 @@ package vip.mate.memory.spi; import io.micrometer.core.instrument.MeterRegistry; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Component; +import vip.mate.agent.context.TokenEstimator; import vip.mate.memory.MemoryProperties; import vip.mate.memory.spi.decorator.MetricsMemoryProvider; import vip.mate.memory.spi.decorator.RetryableMemoryProvider; @@ -81,13 +82,41 @@ public class MemoryManager { * Called once at agent build time (snapshot frozen for session). */ public String buildSystemPromptBlock(Long agentId) { + return buildSystemPromptBlock(agentId, Integer.MAX_VALUE); + } + + /** + * Budgeted variant: providers keep their own per-block caps, but the + * combined output additionally may not exceed {@code budgetTokens} + * (estimated). Provider order is priority order — once the budget is + * spent, later providers are dropped whole and a partially fitting block + * is truncated at a line boundary. Small local context windows need this: + * the individual caps are sized for large cloud models and stack up past + * an 8k/16k window on their own. + */ + public String buildSystemPromptBlock(Long agentId, int budgetTokens) { List blocks = new ArrayList<>(); + int usedTokens = 0; for (MemoryProvider provider : providers) { try { String block = provider.systemPromptBlock(agentId); - if (block != null && !block.isBlank()) { - blocks.add(block); + if (block == null || block.isBlank()) { + continue; } + int blockTokens = TokenEstimator.estimateTokens(block); + if (usedTokens + blockTokens <= budgetTokens) { + blocks.add(block); + usedTokens += blockTokens; + continue; + } + int remaining = budgetTokens - usedTokens; + String truncated = truncateToTokenBudget(block, remaining); + if (!truncated.isBlank()) { + blocks.add(truncated + "\n\n[memory truncated to fit the model context window]"); + } + log.info("[MemoryManager] Memory block budget {} tokens reached at provider '{}' — " + + "remaining providers dropped from the system prompt", budgetTokens, provider.id()); + break; } catch (Exception e) { log.warn("[MemoryManager] Provider '{}' systemPromptBlock() failed: {}", provider.id(), e.getMessage()); @@ -96,6 +125,27 @@ public class MemoryManager { return String.join("\n\n", blocks); } + /** Trim to the last full line that fits the token budget; empty when nothing fits. */ + private static String truncateToTokenBudget(String block, int budgetTokens) { + if (budgetTokens <= 0) { + return ""; + } + StringBuilder kept = new StringBuilder(); + int usedTokens = 0; + for (String line : block.split("\n", -1)) { + int lineTokens = TokenEstimator.estimateTokens(line) + 1; + if (usedTokens + lineTokens > budgetTokens) { + break; + } + if (kept.length() > 0) { + kept.append('\n'); + } + kept.append(line); + usedTokens += lineTokens; + } + return kept.toString(); + } + // ==================== Prefetch / Recall ==================== /** diff --git a/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java b/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java index 4df4700a..29965c7b 100644 --- a/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java +++ b/mateclaw-server/src/main/java/vip/mate/wiki/service/WikiContextService.java @@ -3,6 +3,7 @@ package vip.mate.wiki.service; import lombok.RequiredArgsConstructor; import lombok.extern.slf4j.Slf4j; import org.springframework.stereotype.Service; +import vip.mate.agent.context.TokenEstimator; import vip.mate.wiki.WikiProperties; import vip.mate.wiki.dto.PageSearchResult; import vip.mate.wiki.model.WikiKnowledgeBaseEntity; @@ -47,9 +48,22 @@ public class WikiContextService { * returns snippet + reason instead of just summary. */ public String buildRelevantContext(Long agentId, String userMessage) { + return buildRelevantContext(agentId, userMessage, null); + } + + /** + * Budgeted variant: in addition to the absolute {@code maxContextChars} + * cap, the injected block may not exceed {@code budgetTokens} (estimated). + * Null budget keeps the previous chars-only behavior. Small local context + * windows need this — the chars cap is sized for large cloud models. + */ + public String buildRelevantContext(Long agentId, String userMessage, Integer budgetTokens) { if (!properties.isEnabled() || userMessage == null || userMessage.isBlank()) { return ""; } + if (budgetTokens != null && budgetTokens <= 0) { + return ""; + } // Skip retrieval for continuation / acknowledgement turns — observed // in production: a user reply of "继续" produced top-5 hits dominated @@ -94,15 +108,26 @@ public class WikiContextService { "e.g. 「来源：[[页面标题]]」or「(来源：页面标题)」.]\n\n"); int totalChars = 0; int maxChars = properties.getMaxContextChars(); + int totalTokens = TokenEstimator.estimateTokens(sb.toString()); + boolean anyEntry = false; for (PageSearchResult hit : hits) { String entry = buildContextEntry(hit); - if (totalChars + entry.length() > maxChars) { + int entryTokens = TokenEstimator.estimateTokens(entry); + if (totalChars + entry.length() > maxChars + || (budgetTokens != null && totalTokens + entryTokens > budgetTokens)) { sb.append("- ... (use wiki_search_pages for more)\n"); break; } sb.append(entry); totalChars += entry.length(); + totalTokens += entryTokens; + anyEntry = true; + } + // A budget too small for even one entry yields a header-only block — + // pure overhead. Skip the injection entirely; wiki tools stay usable. + if (!anyEntry && budgetTokens != null) { + return ""; } sb.append(""); return sb.toString(); diff --git a/mateclaw-server/src/test/java/vip/mate/agent/context/PrefixBudgetPlannerTest.java b/mateclaw-server/src/test/java/vip/mate/agent/context/PrefixBudgetPlannerTest.java new file mode 100644 index 00000000..21ae3a87 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/context/PrefixBudgetPlannerTest.java @@ -0,0 +1,121 @@ +package vip.mate.agent.context; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import vip.mate.config.ConversationWindowProperties; +import vip.mate.config.PrefixBudgetProperties; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Unit tests for {@link PrefixBudgetPlanner} — profile selection, share + * allocation, and the disabled/rollback path. + */ +class PrefixBudgetPlannerTest { + + private PrefixBudgetProperties properties; + private ConversationWindowProperties windowProperties; + private PrefixBudgetPlanner planner; + + @BeforeEach + void setUp() { + properties = new PrefixBudgetProperties(); + windowProperties = new ConversationWindowProperties(); + planner = new PrefixBudgetPlanner(properties, windowProperties); + } + + @Test + @DisplayName("large window → NORMAL profile, budget = max*ratio − base − tools") + void normalProfile() { + PrefixBudgetPlan plan = planner.plan(128000, 2000, 3000); + assertEquals(PrefixBudgetPlan.Profile.NORMAL, plan.profile()); + assertEquals(128000, plan.effectiveMaxTokens()); + assertEquals((int) (128000 * 0.35) - 2000 - 3000, plan.injectionBudgetTokens()); + assertTrue(plan.enabled()); + } + + @Test + @DisplayName("16k window → COMPACT profile with tightened ratio") + void compactProfile() { + PrefixBudgetPlan plan = planner.plan(16384, 1000, 1000); + assertEquals(PrefixBudgetPlan.Profile.COMPACT, plan.profile()); + assertEquals((int) (16384 * 0.25) - 2000, plan.injectionBudgetTokens()); + } + + @Test + @DisplayName("4k window → MINIMAL profile") + void minimalProfile() { + PrefixBudgetPlan plan = planner.plan(4096, 500, 500); + assertEquals(PrefixBudgetPlan.Profile.MINIMAL, plan.profile()); + assertEquals(Math.max(0, (int) (4096 * 0.15) - 1000), plan.injectionBudgetTokens()); + } + + @Test + @DisplayName("oversized base prompt + tools clamp the budget to zero, never negative") + void budgetNeverNegative() { + PrefixBudgetPlan plan = planner.plan(8192, 9000, 5000); + assertEquals(0, plan.injectionBudgetTokens()); + assertEquals(0, plan.memoryTokens()); + assertEquals(0, plan.wikiTokens()); + } + + @Test + @DisplayName("shares split the injection budget and sum to at most the budget") + void sharesSplitBudget() { + PrefixBudgetPlan plan = planner.plan(128000, 0, 0); + int budget = plan.injectionBudgetTokens(); + // ±1 token tolerance: share division is floating point. + assertEquals(budget * 0.35, plan.memoryTokens(), 1.0); + assertEquals(budget * 0.30, plan.wikiTokens(), 1.0); + assertEquals(budget * 0.20, plan.skillCatalogTokens(), 1.0); + assertEquals(budget * 0.10, plan.extensionCatalogTokens(), 1.0); + assertEquals(budget * 0.05, plan.ledgerTokens(), 1.0); + int sum = plan.memoryTokens() + plan.wikiTokens() + plan.skillCatalogTokens() + + plan.extensionCatalogTokens() + plan.ledgerTokens(); + assertTrue(sum <= budget); + } + + @Test + @DisplayName("shares not summing to 1 are normalized") + void sharesNormalized() { + properties.getShares().setMemory(2.0); + properties.getShares().setWiki(2.0); + properties.getShares().setSkill(0.0); + properties.getShares().setExtensionCatalog(0.0); + properties.getShares().setLedger(0.0); + PrefixBudgetPlan plan = planner.plan(128000, 0, 0); + assertEquals(plan.injectionBudgetTokens() / 2, plan.memoryTokens()); + assertEquals(plan.injectionBudgetTokens() / 2, plan.wikiTokens()); + assertEquals(0, plan.skillCatalogTokens()); + } + + @Test + @DisplayName("disabled → unlimited plan, previous behavior") + void disabledYieldsUnlimited() { + properties.setEnabled(false); + PrefixBudgetPlan plan = planner.plan(8192, 100, 100); + assertFalse(plan.enabled()); + assertEquals(Integer.MAX_VALUE, plan.memoryTokens()); + assertEquals(Integer.MAX_VALUE, plan.wikiTokens()); + assertEquals(8192, plan.effectiveMaxTokens()); + } + + @Test + @DisplayName("null/zero effective window falls back to the global default") + void nullWindowFallsBackToGlobalDefault() { + PrefixBudgetPlan plan = planner.plan(null, 0, 0); + assertEquals(windowProperties.getDefaultMaxInputTokens(), plan.effectiveMaxTokens()); + assertEquals(PrefixBudgetPlan.Profile.NORMAL, plan.profile()); + } + + @Test + @DisplayName("small windows compact later (0.85), large windows keep the default ratio") + void compactTriggerRatioAdapts() { + assertEquals(0.75, planner.compactTriggerRatioFor(128000, 0.75)); + assertEquals(0.85, planner.compactTriggerRatioFor(16384, 0.75)); + assertEquals(0.85, planner.compactTriggerRatioFor(4096, 0.75)); + properties.setEnabled(false); + assertEquals(0.75, planner.compactTriggerRatioFor(4096, 0.75)); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/memory/MemoryManagerBudgetTest.java b/mateclaw-server/src/test/java/vip/mate/memory/MemoryManagerBudgetTest.java new file mode 100644 index 00000000..3c042afb --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/memory/MemoryManagerBudgetTest.java @@ -0,0 +1,96 @@ +package vip.mate.memory; + +import io.micrometer.core.instrument.MeterRegistry; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.ObjectProvider; +import vip.mate.memory.spi.MemoryManager; +import vip.mate.memory.spi.MemoryProvider; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Tests for the token-budgeted system-prompt block assembly in + * {@link MemoryManager} — provider order is priority order, later providers + * drop whole once the budget is spent, a partially fitting block truncates + * at a line boundary. + */ +class MemoryManagerBudgetTest { + + /** CJK chars estimate ≈ 1 token each, which makes budgets easy to reason about. */ + private static final String BLOCK_A = "甲".repeat(300); + private static final String BLOCK_B = "乙".repeat(300); + + private static MemoryProvider provider(String id, String block) { + return new MemoryProvider() { + @Override + public String id() { + return id; + } + + @Override + public String systemPromptBlock(Long agentId) { + return block; + } + }; + } + + private static MemoryManager manager(MemoryProvider... providers) { + ObjectProvider noRegistry = new ObjectProvider<>() { + @Override + public MeterRegistry getObject(Object... args) { + throw new UnsupportedOperationException(); + } + + @Override + public MeterRegistry getIfAvailable() { + return null; + } + }; + return new MemoryManager(List.of(providers), new MemoryProperties(), noRegistry); + } + + @Test + @DisplayName("unbudgeted call joins every provider block — previous behavior") + void unbudgetedJoinsAll() { + MemoryManager manager = manager(provider("a", BLOCK_A), provider("b", BLOCK_B)); + String result = manager.buildSystemPromptBlock(1L); + assertTrue(result.contains(BLOCK_A)); + assertTrue(result.contains(BLOCK_B)); + } + + @Test + @DisplayName("budget exhausted after the first block → later provider dropped whole") + void budgetDropsLaterProviders() { + MemoryManager manager = manager(provider("a", BLOCK_A), provider("b", BLOCK_B)); + // 300 CJK chars ≈ 300 tokens; 350 fits block A but not A+B. + String result = manager.buildSystemPromptBlock(1L, 350); + assertTrue(result.contains(BLOCK_A)); + assertFalse(result.contains("乙")); + } + + @Test + @DisplayName("single multi-line block over budget truncates at a line boundary with a marker") + void oversizedBlockTruncatesAtLineBoundary() { + String multiLine = ("行".repeat(100) + "\n").repeat(10).trim(); + MemoryManager manager = manager(provider("a", multiLine)); + String result = manager.buildSystemPromptBlock(1L, 350); + assertTrue(result.contains("[memory truncated to fit the model context window]")); + // Only whole lines are kept — every kept content line is the full 100-char line. + for (String line : result.split("\n")) { + if (line.startsWith("行")) { + assertEquals(100, line.length()); + } + } + assertTrue(result.length() < multiLine.length()); + } + + @Test + @DisplayName("zero budget yields an empty block, not an exception") + void zeroBudgetYieldsEmpty() { + MemoryManager manager = manager(provider("a", BLOCK_A)); + assertEquals("", manager.buildSystemPromptBlock(1L, 0)); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java b/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java new file mode 100644 index 00000000..e5047a6e --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/wiki/service/WikiContextServiceBudgetTest.java @@ -0,0 +1,87 @@ +package vip.mate.wiki.service; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; +import vip.mate.wiki.WikiProperties; +import vip.mate.wiki.dto.PageSearchResult; +import vip.mate.wiki.model.WikiKnowledgeBaseEntity; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; + +/** + * Tests for the token-budgeted knowledge-base relevance injection in + * {@link WikiContextService}. + */ +class WikiContextServiceBudgetTest { + + private WikiKnowledgeBaseService kbService; + private HybridRetriever hybridRetriever; + private WikiProperties properties; + private WikiContextService service; + + @BeforeEach + void setUp() { + kbService = Mockito.mock(WikiKnowledgeBaseService.class); + hybridRetriever = Mockito.mock(HybridRetriever.class); + WikiPageService pageService = Mockito.mock(WikiPageService.class); + properties = new WikiProperties(); + service = new WikiContextService(kbService, pageService, hybridRetriever, properties); + + WikiKnowledgeBaseEntity kb = new WikiKnowledgeBaseEntity(); + kb.setId(7L); + Mockito.when(kbService.resolvePrimaryKb(anyLong())).thenReturn(kb); + } + + private void stubHits(int count, int excerptChars) { + List hits = new java.util.ArrayList<>(); + for (int i = 0; i < count; i++) { + hits.add(PageSearchResult.of("page-" + i, "页面" + i, null, + "摘".repeat(excerptChars), List.of("keyword"), null, 1.0)); + } + Mockito.when(hybridRetriever.search(any(), anyString(), anyString(), anyInt())) + .thenReturn(hits); + } + + @Test + @DisplayName("null budget keeps the chars-only behavior — all hits injected") + void nullBudgetKeepsAllHits() { + stubHits(3, 200); + String result = service.buildRelevantContext(1L, "如何配置数据库连接", null); + assertTrue(result.contains("page-0")); + assertTrue(result.contains("page-2")); + } + + @Test + @DisplayName("token budget cuts the tail hits and appends the search hint") + void budgetCutsTailHits() { + stubHits(3, 400); // each entry ≈ 400+ tokens (CJK) + String result = service.buildRelevantContext(1L, "如何配置数据库连接", 600); + assertTrue(result.contains("page-0")); + assertFalse(result.contains("page-2")); + assertTrue(result.contains("use wiki_search_pages for more")); + } + + @Test + @DisplayName("budget too small for even one entry → injection skipped entirely") + void tinyBudgetSkipsInjection() { + stubHits(3, 400); + String result = service.buildRelevantContext(1L, "如何配置数据库连接", 50); + assertEquals("", result); + } + + @Test + @DisplayName("zero or negative budget short-circuits without retrieval") + void zeroBudgetShortCircuits() { + String result = service.buildRelevantContext(1L, "如何配置数据库连接", 0); + assertEquals("", result); + Mockito.verifyNoInteractions(hybridRetriever); + } +}