diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/ConversationWindowManager.java b/mateclaw-server/src/main/java/vip/mate/agent/context/ConversationWindowManager.java index 26de0ac8..345ba438 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/context/ConversationWindowManager.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/ConversationWindowManager.java @@ -170,6 +170,15 @@ public class ConversationWindowManager { * conversation in a compaction storm. */ private final ConcurrentHashMap ptlForceCompactAt = new ConcurrentHashMap<>(); + /** + * Default max input tokens for the configured model window. Surfaced for + * the per-loop budgeter so the L1 (multi-turn compaction) and L2 + * (per-iteration trim) layers stay calibrated to the same number. + */ + public int getDefaultMaxInputTokens() { + return properties != null ? properties.getDefaultMaxInputTokens() : 0; + } + /** Cooldown window after a structured PTL compaction during which a * follow-up PTL is downgraded to tail-only. Picked so a single ReAct * loop that retries within seconds can't burn another summary LLM diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/LoopBudgetConfig.java b/mateclaw-server/src/main/java/vip/mate/agent/context/LoopBudgetConfig.java new file mode 100644 index 00000000..01228d6d --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/LoopBudgetConfig.java @@ -0,0 +1,140 @@ +package vip.mate.agent.context; + +/** + * Configuration for per-reasoning-loop message budgeting. + * + *

Used by {@link LoopMessageBudgeter} to decide when and how to trim the + * working message list that a ReAct iteration hands to the LLM. Distinct from + * the multi-turn history compression configured by + * {@link vip.mate.config.ConversationWindowProperties}: this one applies inside + * a single user turn while the ReAct loop accumulates reasoning steps and + * tool-call/tool-response pairs. + * + *

Field semantics: + *

+ */ +public record LoopBudgetConfig( + int triggerTokens, + int keepTailTokens, + int minTailMessages, + double tailSoftCeilingRatio, + int reservedPrefixTokens, + int targetMaxMessages) { + + /** Smallest useful trigger threshold; below this budgeting is effectively disabled. */ + public static final int MIN_TRIGGER_TOKENS = 1_000; + + /** Smallest sensible tail budget; below this even one observation may not fit. */ + public static final int MIN_TAIL_TOKENS = 2_000; + + /** Floor on minTailMessages — fewer than 3 collapses recent context too aggressively. */ + public static final int MIN_TAIL_MESSAGES_FLOOR = 3; + + /** Floor on the soft ceiling ratio — anything below 1.0 is degenerate. */ + public static final double MIN_TAIL_SOFT_CEILING_RATIO = 1.0; + + /** Smallest sensible target cap; below this even a normal ReAct loop trips it. */ + public static final int MIN_TARGET_MAX = 20; + + public LoopBudgetConfig { + if (triggerTokens < MIN_TRIGGER_TOKENS) { + throw new IllegalArgumentException( + "triggerTokens must be >= " + MIN_TRIGGER_TOKENS + ", got " + triggerTokens); + } + if (keepTailTokens < MIN_TAIL_TOKENS) { + throw new IllegalArgumentException( + "keepTailTokens must be >= " + MIN_TAIL_TOKENS + ", got " + keepTailTokens); + } + if (minTailMessages < MIN_TAIL_MESSAGES_FLOOR) { + throw new IllegalArgumentException( + "minTailMessages must be >= " + MIN_TAIL_MESSAGES_FLOOR + + ", got " + minTailMessages); + } + if (tailSoftCeilingRatio < MIN_TAIL_SOFT_CEILING_RATIO) { + throw new IllegalArgumentException( + "tailSoftCeilingRatio must be >= " + MIN_TAIL_SOFT_CEILING_RATIO + + ", got " + tailSoftCeilingRatio); + } + if (reservedPrefixTokens < 0) { + throw new IllegalArgumentException( + "reservedPrefixTokens must be >= 0, got " + reservedPrefixTokens); + } + if (targetMaxMessages < MIN_TARGET_MAX) { + throw new IllegalArgumentException( + "targetMaxMessages must be >= " + MIN_TARGET_MAX + + ", got " + targetMaxMessages); + } + if (keepTailTokens >= triggerTokens) { + throw new IllegalArgumentException( + "keepTailTokens (" + keepTailTokens + ") must be < triggerTokens (" + + triggerTokens + ") — otherwise budgeting would never reduce anything"); + } + } + + /** Tail budget after applying the soft ceiling. */ + public int tailSoftCeilingTokens() { + return (int) (keepTailTokens * tailSoftCeilingRatio); + } + + /** + * Derive a sensible config from a model's context window. The ratios were + * chosen so the budgeter triggers well before the model's actual limit and + * leaves enough headroom for the LLM's own response. + * + * + */ + public static LoopBudgetConfig forContext(int contextWindowTokens) { + if (contextWindowTokens <= 0) { + contextWindowTokens = 32_000; + } + int trigger = Math.max(MIN_TRIGGER_TOKENS, (int) (contextWindowTokens * 0.50)); + int tail = Math.max(MIN_TAIL_TOKENS, (int) (contextWindowTokens * 0.30)); + if (tail >= trigger) { + tail = Math.max(MIN_TAIL_TOKENS, trigger - MIN_TRIGGER_TOKENS); + } + return new LoopBudgetConfig(trigger, tail, 4, 1.5, 0, 200); + } + + /** Return a copy with {@code reservedPrefixTokens} replaced. */ + public LoopBudgetConfig withReservedPrefixTokens(int reservedPrefixTokens) { + return new LoopBudgetConfig(triggerTokens, keepTailTokens, minTailMessages, + tailSoftCeilingRatio, reservedPrefixTokens, targetMaxMessages); + } +} diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/LoopMessageBudgeter.java b/mateclaw-server/src/main/java/vip/mate/agent/context/LoopMessageBudgeter.java new file mode 100644 index 00000000..356ac616 --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/LoopMessageBudgeter.java @@ -0,0 +1,296 @@ +package vip.mate.agent.context; + +import lombok.extern.slf4j.Slf4j; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.stereotype.Component; + +import java.util.ArrayList; +import java.util.List; + +/** + * Per-ReAct-loop message budgeter. Bounds the working message list a Reasoning + * iteration hands to the LLM while preserving five invariants that, when + * violated, either produce off-topic answers or 400s from strict providers: + * + *
    + *
  1. System prompt(s) — all consecutive {@link SystemMessage}s at + * the head stay verbatim. Production agents commonly have multiple + * (SOUL, AGENTS, runtime context, wiki, tool prompt, skill catalog).
  2. + *
  3. Turn anchor — the latest {@link UserMessage} is never dropped. + * Stitched in when an aggressive cut would otherwise lose it.
  4. + *
  5. Tool-call/response pair integrity — every assistant tool_call + * reaches the model with its matching tool_response, and vice versa. + * Delegated to {@link ToolPairSanitizer}.
  6. + *
  7. Token budget over message count — tail sized by token estimate + * so a small ReAct loop with fat observations and a large loop with + * thin observations both fit one config.
  8. + *
  9. Minimum tail messages — at least {@code minTailMessages} + * entries survive even when a single message is bigger than the + * hard tail budget. Prevents collapsing recent reasoning to one row + * when the latest tool output is huge.
  10. + *
+ * + *

The trigger threshold compares {@code historyTokens + + * reservedPrefixTokens} against {@code triggerTokens}; this keeps the + * budgeter calibrated against the entire prompt the LLM will see, not just + * the message list (the L1 compactor uses the same arithmetic). + * + *

Distinct from {@link ConversationWindowManager}: that one runs once per + * user turn and produces a structured LLM summary for the accumulated + * multi-turn history. This one runs per reasoning iteration on top of + * whatever {@code ConversationWindowManager} already produced, bounding the + * intra-turn ReAct accumulation. + * + *

Stateless and side-effect-free for callers; safe to call from + * concurrent reasoning threads. The orphan-removal pass mutates a freshly + * allocated local list, never the caller's input. + */ +@Slf4j +@Component +public class LoopMessageBudgeter { + + /** Outcome of a budgeting pass. */ + public record Result(List messages, BudgetTrace trace) {} + + /** + * Structured trace of a single budgeting decision. All counts and token + * figures refer to {@link Message} entries. + * + *

+ */ + public record BudgetTrace( + int originalCount, + int originalTokens, + int finalCount, + int finalTokens, + int reservedPrefixTokens, + int headKept, + int tailKept, + int droppedMiddle, + int orphansRemoved, + boolean anchorEnforced, + boolean anchorStitched, + boolean targetMaxTripped, + boolean capExceededForPairIntegrity, + boolean minTailFloorApplied, + boolean triggered, + boolean modified) { + + /** Trace for the no-op case (budget not triggered). */ + public static BudgetTrace untouched(int count, int tokens, int prefixTokens, int headKept) { + return new BudgetTrace(count, tokens, count, tokens, prefixTokens, headKept, + count - headKept, 0, 0, + false, false, false, false, false, false, false); + } + } + + /** Apply the loop budget to {@code messages}. Pure function; never mutates the input. */ + public Result budget(List messages, LoopBudgetConfig cfg) { + if (messages == null || messages.isEmpty()) { + return new Result(messages == null ? List.of() : messages, + BudgetTrace.untouched(0, 0, cfg.reservedPrefixTokens(), 0)); + } + int originalCount = messages.size(); + int historyTokens = TokenEstimator.estimateTokens(messages); + int headEnd = findHeadEnd(messages); + + // Budget against the full prompt (history + prefix), so the trigger + // matches what the LLM would actually receive — not just the + // history slice. Prefix covers system prompt, skill catalog, + // runtime context, wiki, tool schemas, output reserve. + int promptTokens = historyTokens + cfg.reservedPrefixTokens(); + + // Below both thresholds → forward unchanged. + if (promptTokens < cfg.triggerTokens() && originalCount < cfg.targetMaxMessages()) { + return new Result(messages, + BudgetTrace.untouched(originalCount, historyTokens, + cfg.reservedPrefixTokens(), headEnd)); + } + + // 1. Token-budgeted tail cut. Walk backward from the end; the + // earliest index whose suffix fits within keepTailTokens is the + // proposed boundary. + int hardTailStart = findTailCutByTokens(messages, headEnd, cfg.keepTailTokens()); + + // 2. Min-tail floor: if the hard cut keeps fewer than minTailMessages, + // pull back to keep at least that many — but only up to the soft + // ceiling. Without this, one giant tool output can collapse the + // tail to a single row and lose recent reasoning context. + boolean minTailFloorApplied = false; + int tailStart = hardTailStart; + int hardTailCount = originalCount - hardTailStart; + if (hardTailCount < cfg.minTailMessages()) { + int floorTailStart = Math.max(headEnd, originalCount - cfg.minTailMessages()); + // Honor the soft ceiling: if even the floor count would consume + // more than tailSoftCeilingTokens, accept it (the floor wins, + // since the alternative is losing recent reasoning entirely). + tailStart = floorTailStart; + minTailFloorApplied = true; + } else { + // Apply the soft ceiling: if the hard cut undershoots the soft + // ceiling (i.e. there's slack), keep going. We already cut to + // the hard budget so there's no need to expand here — the soft + // ceiling acts as a guard rail for the floor/pull-back path, + // not as a relaxation of the normal cut. + } + + // 3. Anchor: never drop the latest UserMessage. Pull tail back if + // needed (cheap — just moves the boundary). + boolean anchorEnforced = false; + int anchorIdx = findLatestUserMessageIdx(messages, headEnd); + if (anchorIdx >= 0 && anchorIdx < tailStart) { + tailStart = anchorIdx; + anchorEnforced = true; + } + + // 4. Tool-pair integrity at the boundary: if tailStart sits inside a + // tool pair, pull back so the pair survives whole. Delegated to + // the shared sanitizer. + int beforePairPullBack = tailStart; + tailStart = ToolPairSanitizer.pullBackToToolPairBoundary(messages, headEnd, tailStart); + + // 5. Target max safety net. The pair-integrity pull-back may have + // pushed final count above the soft cap; we re-evaluate and try + // to enforce, but pair integrity wins over count cap. + boolean targetMaxTripped = false; + boolean anchorStitched = false; + boolean capExceededForPairIntegrity = false; + int targetTailCap = Math.max(0, cfg.targetMaxMessages() - headEnd); + if (targetTailCap > 0 && (originalCount - tailStart) > targetTailCap) { + int provisionalTailStart = originalCount - targetTailCap; + boolean stitchNeeded = anchorIdx >= 0 && anchorIdx < provisionalTailStart; + int reservedForStitchedAnchor = stitchNeeded ? 1 : 0; + int recentTailCap = Math.max(1, targetTailCap - reservedForStitchedAnchor); + int newTailStart = originalCount - recentTailCap; + int adjustedTailStart = ToolPairSanitizer.pullBackToToolPairBoundary( + messages, headEnd, newTailStart); + if (adjustedTailStart < newTailStart) { + // Pair integrity prevailed over the cap; honestly record that + // the final count will exceed targetMaxMessages. + capExceededForPairIntegrity = true; + } + tailStart = adjustedTailStart; + targetMaxTripped = true; + anchorStitched = anchorIdx >= 0 && anchorIdx < tailStart; + } + + // Detect anchor stitching from the tool-pair pull-back path too: + // pull-back may have moved tailStart earlier than the anchor index + // (rare, but possible if the pair anchor is in the head section). + if (!anchorStitched && anchorIdx >= 0 && anchorIdx < tailStart) { + anchorStitched = true; + } + + // 6. Build the trimmed list: head + [stitched anchor?] + tail. + int estimated = headEnd + (anchorStitched ? 1 : 0) + (originalCount - tailStart); + List trimmed = new ArrayList<>(estimated); + trimmed.addAll(messages.subList(0, headEnd)); + if (anchorStitched) { + trimmed.add(messages.get(anchorIdx)); + } + trimmed.addAll(messages.subList(tailStart, originalCount)); + + // 7. Tool-pair invariant: cross-boundary orphans cleaned up. The + // pull-back at step 4 handles the boundary case but a head-section + // Assistant(tool_calls) whose responses fell in the dropped middle + // still needs the bidirectional pass. + int orphans = ToolPairSanitizer.removeOrphans(trimmed); + + int finalCount = trimmed.size(); + int finalTokens = TokenEstimator.estimateTokens(trimmed); + int droppedMiddle = originalCount - finalCount; + + // Touch the unused locals so the compiler doesn't warn — they're + // useful in the trace's narrative but the actual cut already + // committed. + if (beforePairPullBack != tailStart) { + // pair pull-back moved the boundary; logged via trace fields + } + + boolean modified = (finalCount != originalCount) || (orphans > 0); + + return new Result(trimmed, new BudgetTrace( + originalCount, historyTokens, + finalCount, finalTokens, + cfg.reservedPrefixTokens(), + headEnd, + finalCount - headEnd, + droppedMiddle, + orphans, + anchorEnforced, + anchorStitched, + targetMaxTripped, + capExceededForPairIntegrity, + minTailFloorApplied, + /* triggered */ true, + modified)); + } + + // ------------------------------------------------------------------------ + // Internals + // ------------------------------------------------------------------------ + + private static int findHeadEnd(List messages) { + int i = 0; + while (i < messages.size() && messages.get(i) instanceof SystemMessage) { + i++; + } + return i; + } + + /** + * Walk backward from the end accumulating per-message token estimates. + * Return the earliest index whose suffix fits within {@code keepTokens}. + * Always returns a value in {@code [headEnd, messages.size())} so the + * tail is non-empty. + */ + private static int findTailCutByTokens(List messages, int headEnd, int keepTokens) { + int n = messages.size(); + if (n <= headEnd) { + return n; + } + int acc = 0; + for (int i = n - 1; i >= headEnd; i--) { + int t = TokenEstimator.estimateTokens(messages.get(i)); + if (acc + t > keepTokens && i < n - 1) { + return i + 1; + } + acc += t; + } + return headEnd; + } + + /** Index of the latest {@link UserMessage} at or after {@code headEnd}; -1 if none. */ + private static int findLatestUserMessageIdx(List messages, int headEnd) { + for (int i = messages.size() - 1; i >= headEnd; i--) { + if (messages.get(i) instanceof UserMessage) { + return i; + } + } + return -1; + } +} diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/ToolPairSanitizer.java b/mateclaw-server/src/main/java/vip/mate/agent/context/ToolPairSanitizer.java new file mode 100644 index 00000000..ac3ac94f --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/ToolPairSanitizer.java @@ -0,0 +1,192 @@ +package vip.mate.agent.context; + +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.ToolResponseMessage; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * Pure-function utilities that enforce the OpenAI-compatible + * tool_call ↔ tool_response pairing invariant: + * + *
    + *
  • Every {@code tool_call.id} on an {@link AssistantMessage} has a + * matching {@code tool_response.id} on a {@link ToolResponseMessage} + * after it in the list.
  • + *
  • Every {@code tool_response.id} on a {@link ToolResponseMessage} has + * a matching {@code tool_call.id} on an {@link AssistantMessage} + * before it.
  • + *
  • No empty/null ids on either side.
  • + *
+ * + *

Violating either rule causes strict providers (kimi-code, anthropic in + * tool-use mode, OpenAI's responses API on certain models) to reject the + * request with a 400 error such as + * {@code "tool_call_id is not found"}. This sanitizer is the single source of + * truth for that invariant — any trim / cut / window logic should run its + * pre/post passes here rather than reimplementing them. + * + *

All methods are {@code static} and side-effect-free except where + * documented (e.g. {@link #removeOrphans(List)} mutates the list in place to + * avoid an extra allocation hot in the reasoning loop). They never touch the + * input list when no fix is needed. + */ +public final class ToolPairSanitizer { + + private ToolPairSanitizer() { + // utility class + } + + /** + * Pull a proposed cut boundary earlier so an Assistant(tool_calls) that + * issued ids matching {@link ToolResponseMessage}s in the kept tail + * survives into the tail alongside its responses. Prevents producing an + * orphan response at the cut boundary in the first place. + * + * @param messages full message list (read-only) + * @param headEnd index after the last protected head message + * @param tailStart proposed boundary; messages at and after this index + * are kept, those between {@code headEnd} and + * {@code tailStart} are dropped + * @return possibly-earlier {@code tailStart} that keeps tool pairs whole + */ + public static int pullBackToToolPairBoundary(List messages, int headEnd, int tailStart) { + if (tailStart <= headEnd || messages == null || messages.isEmpty()) { + return tailStart; + } + Set tailResponseIds = new HashSet<>(); + for (int i = tailStart; i < messages.size(); i++) { + if (messages.get(i) instanceof ToolResponseMessage trm) { + for (ToolResponseMessage.ToolResponse r : trm.getResponses()) { + if (r.id() != null && !r.id().isEmpty()) { + tailResponseIds.add(r.id()); + } + } + } + } + if (tailResponseIds.isEmpty()) { + return tailStart; + } + for (int i = tailStart - 1; i >= headEnd; i--) { + Message m = messages.get(i); + if (m instanceof AssistantMessage am && am.getToolCalls() != null) { + boolean overlaps = am.getToolCalls().stream() + .anyMatch(tc -> tc.id() != null && tailResponseIds.contains(tc.id())); + if (overlaps) { + return i; + } + } + } + return tailStart; + } + + /** + * Iteratively remove tool-pair orphans from {@code messages} (mutates the + * list in place). Two shapes are handled: + * + *

P0: a {@link ToolResponseMessage} whose response id has no + * matching assistant tool_call in the list. + * + *

P1: an {@link AssistantMessage} whose every tool_call id + * has no matching response in the list. (An assistant with both matched + * and unmatched calls is left alone — removing it would harm more than + * it helps; strict providers tolerate extra calls more readily than + * dropping the whole assistant message.) + * + *

Iterates until convergence: removing an assistant for P1 can expose + * a P0 orphan that needs cleaning, and vice versa. + * + *

Also removes any tool_call or tool_response with a null or empty id + * — those have no useful pairing semantics and confuse both the strict + * providers and the matching logic. + * + * @return total number of messages removed across all passes + */ + public static int removeOrphans(List messages) { + if (messages == null || messages.isEmpty()) { + return 0; + } + int totalRemoved = 0; + boolean changed; + do { + Set callIds = new HashSet<>(); + Set respIds = new HashSet<>(); + for (Message m : messages) { + if (m instanceof AssistantMessage am && am.getToolCalls() != null) { + for (AssistantMessage.ToolCall tc : am.getToolCalls()) { + if (tc.id() != null && !tc.id().isEmpty()) { + callIds.add(tc.id()); + } + } + } + if (m instanceof ToolResponseMessage trm) { + for (ToolResponseMessage.ToolResponse r : trm.getResponses()) { + if (r.id() != null && !r.id().isEmpty()) { + respIds.add(r.id()); + } + } + } + } + int before = messages.size(); + messages.removeIf(m -> { + if (m instanceof ToolResponseMessage trm) { + // P0: a response with a null/empty id, or whose id has + // no matching tool_call. + return trm.getResponses().stream().anyMatch(r -> + r.id() == null || r.id().isEmpty() || !callIds.contains(r.id())); + } + if (m instanceof AssistantMessage am && am.getToolCalls() != null + && !am.getToolCalls().isEmpty()) { + // P1: every tool_call on this assistant has no matching response. + return am.getToolCalls().stream().allMatch(tc -> + tc.id() == null || tc.id().isEmpty() || !respIds.contains(tc.id())); + } + return false; + }); + int removed = before - messages.size(); + totalRemoved += removed; + changed = removed > 0; + } while (changed); + return totalRemoved; + } + + /** + * Post-condition check: returns {@code true} iff {@code messages} + * satisfies the pairing invariant — every assistant tool_call has a + * matching response after it, every response has a matching call before + * it, all ids are non-empty. Intended for tests and defensive asserts; + * production code should run {@link #removeOrphans(List)} which + * guarantees this holds on return. + */ + public static boolean isPaired(List messages) { + if (messages == null || messages.isEmpty()) { + return true; + } + Set callIds = new HashSet<>(); + Set respIds = new HashSet<>(); + for (Message m : messages) { + if (m instanceof AssistantMessage am && am.getToolCalls() != null) { + for (AssistantMessage.ToolCall tc : am.getToolCalls()) { + if (tc.id() == null || tc.id().isEmpty()) return false; + callIds.add(tc.id()); + } + } + if (m instanceof ToolResponseMessage trm) { + for (ToolResponseMessage.ToolResponse r : trm.getResponses()) { + if (r.id() == null || r.id().isEmpty()) return false; + respIds.add(r.id()); + } + } + } + for (String c : callIds) { + if (!respIds.contains(c)) return false; + } + for (String r : respIds) { + if (!callIds.contains(r)) return false; + } + return true; + } +} 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 3ee70aed..9070f18d 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 @@ -21,7 +21,10 @@ import vip.mate.agent.GraphEventPublisher; 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.LoopMessageBudgeter; import vip.mate.agent.context.RuntimeContextInjector; +import vip.mate.agent.context.TokenEstimator; import vip.mate.agent.graph.state.FinishReason; import vip.mate.agent.graph.state.MateClawStateAccessor; import vip.mate.agent.graph.state.MateClawStateKeys; @@ -71,6 +74,36 @@ public class ReasoningNode implements NodeAction { */ private static final int DEFAULT_MAX_OUTPUT_TOKENS = 16384; + /** + * Stateless singleton used to budget the per-iteration working message + * list. Static-final because the budgeter holds no mutable state — the + * choice keeps the existing ReasoningNode constructor surface unchanged + * (it already carries 13 parameters across 5 overloads) and makes the + * dependency obvious to anyone reading the class. + */ + private static final LoopMessageBudgeter LOOP_BUDGETER = new LoopMessageBudgeter(); + + /** + * Fallback context window used when no provider-level value is wired in. + * Calibrated to the same default {@code ConversationWindowProperties} + * uses for its multi-turn budget so the two layers stay in sync. Models + * with smaller windows still benefit — the budgeter triggers earlier on + * raw message volume via {@code absoluteMaxMessages}. + */ + private static final int DEFAULT_LOOP_CONTEXT_WINDOW_TOKENS = 128_000; + + /** + * Conservative buffer added to the per-loop budget's reservedPrefixTokens + * to cover non-history prompt segments that are appended after + * the budget runs: the runtime-rendered skill catalog, runtime-context + * snapshot, wiki injection, progress ledger snapshot, and assorted + * marker SystemMessages. Underestimating here only delays the trigger + * slightly; loop invariants (anchor preservation, tool-pair integrity) + * are unaffected. Sized for a typical agent with 20–30 skills and + * moderate wiki content. + */ + private static final int LOOP_PREFIX_AUXILIARY_RESERVE_TOKENS = 4_000; + /** * DashScope's native chat API caps {@code max_tokens} at 8192 and returns a * 400 {@code InvalidParameter} ("Range of max_tokens should be [1, 8192]") @@ -319,6 +352,20 @@ public class ReasoningNode implements NodeAction { this.progressLedgerService = progressLedgerService; } + /** + * Context window used by the per-loop budgeter. Returns the + * conversation-window manager's effective max input tokens when one is + * wired in (so L1 and L2 stay calibrated to the same model window), + * otherwise the documented fallback. + */ + private int loopContextWindowTokens() { + if (conversationWindowManager != null) { + int v = conversationWindowManager.getDefaultMaxInputTokens(); + if (v > 0) return v; + } + return DEFAULT_LOOP_CONTEXT_WINDOW_TOKENS; + } + public ReasoningNode(ChatModel chatModel, AgentToolSet toolSet, String reasoningEffort, NodeStreamingChatHelper streamingHelper, ConversationWindowManager conversationWindowManager) { @@ -414,81 +461,50 @@ public class ReasoningNode implements NodeAction { systemPrompt = systemPrompt + TOOL_USE_ENFORCEMENT; List messages = accessor.messages(); - // Guard against runaway message list growth. + // Per-loop budget: bound the working message list a single Reasoning + // iteration hands to the LLM. The previous fixed head=4 + tail=36 cut + // could lose the latest UserMessage once the ReAct loop accumulated + // tool calls/observations past ~70 messages — the user's question + // fell into the dropped middle, the LLM lost it, and the agent + // answered off-topic. LoopMessageBudgeter anchors the latest + // UserMessage as undroppable, sizes the tail by token budget instead + // of message count, and keeps the same bidirectional tool-pair + // integrity guard the old block already had. The L2 trim here is + // distinct from ConversationWindowManager (L1): L1 runs once per + // user turn and produces an LLM summary for multi-turn history; L2 + // runs per reasoning iteration on what L1 already produced plus + // intra-turn tool-call growth. // - // CRITICAL: a naive head+tail cut can break the OpenAI-compatible protocol invariant - // that requires tool_call / tool_response pairs to be complete: - // - // P0 (originally observed): AssistantMessage(tool_calls) falls into the dropped gap, - // its ToolResponseMessage lands in the kept tail → provider sees an orphaned - // ToolResponseMessage → kimi-code 400 "tool_call_id is not found". - // - // P1 (symmetric): AssistantMessage(tool_calls) is kept in the head at the boundary, - // its ToolResponseMessage falls into the dropped gap → provider sees an assistant - // tool_call with no matching response → also a 400 on strict providers. - // - // Fix: perform the normal cut, then run an iterative bidirectional integrity pass until - // the list is stable: - // • Remove any ToolResponseMessage whose parent AssistantMessage.tool_calls id was - // dropped (P0). - // • Remove any AssistantMessage whose tool_calls have no matching ToolResponseMessage - // (P1). - // Iterate because a P1 removal could expose a new P0 orphan (and vice versa, though that - // is pathological in practice). With ≤40 messages convergence is always fast. - // Dropping incomplete pairs is safe — prior iterations already processed those - // observations; the LLM needs the summary context, not the raw tool I/O. - final int MAX_LOOP_MESSAGES = 40; - if (messages.size() > MAX_LOOP_MESSAGES) { - log.warn("[ReasoningNode] Messages list too large ({} messages), trimming to {} for conversation {}", - messages.size(), MAX_LOOP_MESSAGES, conversationId); - int headKeep = Math.min(4, messages.size()); - int tailKeep = MAX_LOOP_MESSAGES - headKeep; - int tailStart = messages.size() - tailKeep; - - List trimmed = new ArrayList<>(MAX_LOOP_MESSAGES); - trimmed.addAll(messages.subList(0, headKeep)); - trimmed.addAll(messages.subList(tailStart, messages.size())); - - // Iterative bidirectional integrity pass. - int totalRemoved = 0; - boolean changed; - do { - // Snapshot current tool_call ids and response ids. - Set callIds = new java.util.HashSet<>(); - Set respIds = new java.util.HashSet<>(); - for (Message m : trimmed) { - if (m instanceof AssistantMessage am && am.getToolCalls() != null) { - for (AssistantMessage.ToolCall tc : am.getToolCalls()) callIds.add(tc.id()); - } - if (m instanceof ToolResponseMessage trm) { - for (ToolResponseMessage.ToolResponse r : trm.getResponses()) respIds.add(r.id()); - } - } - int before = trimmed.size(); - trimmed.removeIf(m -> { - // P0: ToolResponseMessage whose parent tool_call was dropped - if (m instanceof ToolResponseMessage trm) { - return trm.getResponses().stream().anyMatch(r -> !callIds.contains(r.id())); - } - // P1: AssistantMessage whose tool_call has no ToolResponseMessage - if (m instanceof AssistantMessage am && am.getToolCalls() != null - && !am.getToolCalls().isEmpty()) { - return am.getToolCalls().stream().anyMatch(tc -> !respIds.contains(tc.id())); - } - return false; - }); - int removed = before - trimmed.size(); - totalRemoved += removed; - changed = removed > 0; - } while (changed); - - if (totalRemoved > 0) { - log.warn("[ReasoningNode] Removed {} message(s) with broken tool_call/response pairs " - + "after trim (bidirectional integrity guard), conv={}", totalRemoved, conversationId); - } - - messages = trimmed; + // Reserved prefix tokens cover the non-history portion of the + // prompt the LLM will receive: system prompt (with tool-use + // enforcement already appended), tool schemas, output reserve, + // and a buffer for skill catalog + runtime context + wiki + // injections that are added downstream. Underestimating here only + // delays the trigger slightly — invariants (anchor, pair integrity) + // still hold once budget fires. + int systemTokens = TokenEstimator.estimateTokens(systemPrompt); + int toolsTokens = TokenEstimator.estimateToolsTokens(toolCallbacks); + int loopReservedPrefixTokens = systemTokens + toolsTokens + + maxOutputTokens + LOOP_PREFIX_AUXILIARY_RESERVE_TOKENS; + LoopBudgetConfig loopCfg = LoopBudgetConfig.forContext(loopContextWindowTokens()) + .withReservedPrefixTokens(loopReservedPrefixTokens); + LoopMessageBudgeter.Result budgeted = LOOP_BUDGETER.budget(messages, loopCfg); + // Only log when the budget actually modified the list — a triggered- + // but-no-op pass is normal (history fits comfortably under the tail + // budget) and would otherwise spam logs every iteration. + if (budgeted.trace().modified()) { + LoopMessageBudgeter.BudgetTrace t = budgeted.trace(); + log.warn("[ReasoningNode] Loop budget trim: {} -> {} msgs (history {} -> {} tokens, " + + "prefix~{}), head={}, tail={}, droppedMiddle={}, orphans={}, " + + "anchorEnforced={}, anchorStitched={}, targetMaxTripped={}, " + + "capExceededForPairIntegrity={}, minTailFloorApplied={}, conv={}", + t.originalCount(), t.finalCount(), t.originalTokens(), t.finalTokens(), + t.reservedPrefixTokens(), + t.headKept(), t.tailKept(), t.droppedMiddle(), t.orphansRemoved(), + t.anchorEnforced(), t.anchorStitched(), t.targetMaxTripped(), + t.capExceededForPairIntegrity(), t.minTailFloorApplied(), conversationId); } + messages = budgeted.messages(); String workspaceBasePath = state.value(vip.mate.agent.graph.state.MateClawStateKeys.WORKSPACE_BASE_PATH, ""); String agentIdStr = state.value(MateClawStateKeys.AGENT_ID, ""); diff --git a/mateclaw-server/src/test/java/vip/mate/agent/context/LoopMessageBudgeterTest.java b/mateclaw-server/src/test/java/vip/mate/agent/context/LoopMessageBudgeterTest.java new file mode 100644 index 00000000..116e8d8b --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/context/LoopMessageBudgeterTest.java @@ -0,0 +1,468 @@ +package vip.mate.agent.context; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.messages.UserMessage; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Behaviour of {@link LoopMessageBudgeter} — the per-ReAct-loop trim that + * replaces the previous fixed head-4 / tail-36 cut. + * + *

The single load-bearing invariant is the anchor: the latest + * {@link UserMessage} must always survive into the LLM-bound list, even when + * the token-budgeted tail would have dropped it. Losing that message is what + * made the agent answer "I am an AI assistant" to a question about Qwen 3.7. + */ +class LoopMessageBudgeterTest { + + private final LoopMessageBudgeter budgeter = new LoopMessageBudgeter(); + + // ---- baselines --------------------------------------------------------- + + @Test + @DisplayName("empty list returns untouched") + void empty_noOp() { + LoopMessageBudgeter.Result r = budgeter.budget(List.of(), defaultCfg()); + assertEquals(0, r.messages().size()); + assertFalse(r.trace().modified()); + } + + @Test + @DisplayName("below trigger thresholds: list is forwarded unchanged") + void belowTrigger_passthrough() { + List input = new ArrayList<>(); + input.add(new SystemMessage("you are a tester")); + input.add(new UserMessage("hi")); + input.add(new AssistantMessage("hello")); + + LoopMessageBudgeter.Result r = budgeter.budget(input, defaultCfg()); + + assertSame(input, r.messages(), "untouched fast-path must return the same list reference"); + assertFalse(r.trace().modified()); + assertEquals(3, r.trace().finalCount()); + assertEquals(1, r.trace().headKept()); + } + + // ---- anchor enforcement (the regression we're fixing) ------------------ + + @Test + @DisplayName("anchor: latest UserMessage is never dropped, tail pulled back to keep it") + void anchor_preventsLatestUserMessageDrop() { + // Simulate the Qwen-3.7 failure shape: one system prompt, the user + // question at index 4, then a flood of tool spam that fills the + // entire token-budgeted tail. The latest UserMessage lives at the + // "middle" of the list and a naive token-budget tail cut drops it. + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("system")); // 0 + msgs.add(new UserMessage("old turn")); // 1 + msgs.add(new AssistantMessage("old reply")); // 2 + for (int i = 0; i < 10; i++) { + // 10 fat noise messages between old turn and the new user turn + msgs.add(new AssistantMessage(fat("noise-pre-" + i))); + } + msgs.add(new UserMessage("查下 qwen 3.7")); // anchor — must survive + for (int i = 0; i < 100; i++) { + // 100 fat tool observations after the new user message + msgs.add(new AssistantMessage(fat("tool-obs-" + i))); + } + int anchorIdx = 13; + assertTrue(msgs.get(anchorIdx) instanceof UserMessage + && ((UserMessage) msgs.get(anchorIdx)).getText().contains("qwen 3.7")); + + // Tail budget intentionally tight so the naive cut would drop the anchor. + LoopBudgetConfig cfg = new LoopBudgetConfig(50_000, 10_000, 4, 1.5, 0, 200); + + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + assertTrue(r.trace().modified(), "budget should have triggered"); + assertTrue(r.trace().anchorEnforced(), + "tail cut had to be pulled back to keep the anchor — that's the whole point"); + assertTrue(containsExact(r.messages(), "查下 qwen 3.7"), + "the user's question MUST remain in the LLM-bound message list"); + // Head must still carry the system prompt. + assertTrue(r.messages().get(0) instanceof SystemMessage); + } + + // ---- token-budget tail vs old fixed-count tail ------------------------- + + @Test + @DisplayName("tail is sized by token estimate, not message count") + void tail_sizedByTokens() { + // Tail of mostly tiny messages can grow to many entries; tail of a + // few huge messages stays small. The same config should produce + // tails of very different message counts. + LoopBudgetConfig cfg = new LoopBudgetConfig(50_000, 8_000, 4, 1.5, 0, 500); + + // Case A: 200 tiny assistant messages — many should survive in tail. + List tiny = new ArrayList<>(); + tiny.add(new SystemMessage("sys")); + tiny.add(new UserMessage("anchor")); + for (int i = 0; i < 200; i++) tiny.add(new AssistantMessage("x" + i)); + // Force trigger by adding bulk to original token total. + for (int i = 0; i < 60; i++) tiny.add(new AssistantMessage(fat("bulk" + i))); + int finalTiny = budgeter.budget(tiny, cfg).trace().finalCount(); + + // Case B: only fat messages. + List big = new ArrayList<>(); + big.add(new SystemMessage("sys")); + big.add(new UserMessage("anchor")); + for (int i = 0; i < 60; i++) big.add(new AssistantMessage(fat("big" + i))); + int finalBig = budgeter.budget(big, cfg).trace().finalCount(); + + assertTrue(finalTiny > finalBig, + "tail with tiny messages should keep more entries than tail with fat ones " + + "(tiny=" + finalTiny + ", big=" + finalBig + ")"); + } + + // ---- head detection ---------------------------------------------------- + + @Test + @DisplayName("head: every consecutive SystemMessage is preserved (not just first 4)") + void head_acceptsManySystemMessages() { + List msgs = new ArrayList<>(); + // 6 system messages: SOUL, AGENTS, runtime context, wiki, tool prompt, + // skill catalog — realistic for production agents. + for (int i = 0; i < 6; i++) msgs.add(new SystemMessage("system-" + i)); + msgs.add(new UserMessage("anchor")); + for (int i = 0; i < 80; i++) msgs.add(new AssistantMessage(fat("obs-" + i))); + + LoopBudgetConfig cfg = new LoopBudgetConfig(20_000, 5_000, 4, 1.5, 0, 200); + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + assertTrue(r.trace().triggered(), "token total exceeds trigger; main path must run"); + assertEquals(6, r.trace().headKept(), + "all six system messages must survive — not the legacy hard-coded 4"); + for (int i = 0; i < 6; i++) { + assertTrue(r.messages().get(i) instanceof SystemMessage, + "head slot " + i + " should be SystemMessage"); + } + } + + // ---- tool-pair integrity ---------------------------------------------- + + @Test + @DisplayName("tool-pair: Assistant(tool_calls) and ToolResponseMessage at boundary stay paired") + void toolPair_pulledBackTogether() { + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("anchor")); + // A long head of fat messages that pushes the tail boundary. + for (int i = 0; i < 50; i++) msgs.add(new AssistantMessage(fat("head-" + i))); + // Tool-pair right at the would-be cut boundary. + msgs.add(asst("call-A")); + msgs.add(resp("call-A")); + // Tail of recent assistant chatter after the pair. + for (int i = 0; i < 5; i++) msgs.add(new AssistantMessage("tail-" + i)); + + LoopBudgetConfig cfg = new LoopBudgetConfig(10_000, 3_000, 4, 1.5, 0, 200); + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + // Both members of the pair survive together; never one without the other. + boolean hasCall = r.messages().stream().anyMatch(m -> m instanceof AssistantMessage am + && am.getToolCalls() != null && am.getToolCalls().stream().anyMatch(tc -> "call-A".equals(tc.id()))); + boolean hasResp = r.messages().stream().anyMatch(m -> m instanceof ToolResponseMessage trm + && trm.getResponses().stream().anyMatch(rr -> "call-A".equals(rr.id()))); + assertEquals(hasCall, hasResp, + "tool_call and its response must both survive or both be removed — never one without the other"); + assertEquals(0, r.trace().orphansRemoved(), + "pull-back at boundary should have prevented any orphan from being produced"); + } + + @Test + @DisplayName("integrity invariant: call-X and resp-X either both survive or both are removed") + void toolPair_integrityInvariant() { + // The pull-back at the boundary normally keeps pairs together, and + // the bidirectional orphan pass catches cross-boundary leftovers. + // The invariant we care about is the post-condition: the final list + // never contains a response without its matching call (or vice versa). + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("anchor")); + for (int i = 0; i < 100; i++) msgs.add(new AssistantMessage(fat("mid-" + i))); + msgs.add(asst("call-X")); + msgs.add(new AssistantMessage(fat("tail-fill-1"))); + msgs.add(resp("call-X")); + msgs.add(new AssistantMessage("tail-fill-2")); + + LoopBudgetConfig cfg = new LoopBudgetConfig(10_000, 2_000, 4, 1.5, 0, 200); + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + boolean hasCall = r.messages().stream().anyMatch(m -> m instanceof AssistantMessage am + && am.getToolCalls() != null + && am.getToolCalls().stream().anyMatch(tc -> "call-X".equals(tc.id()))); + boolean hasResp = r.messages().stream().anyMatch(m -> m instanceof ToolResponseMessage trm + && trm.getResponses().stream().anyMatch(rr -> "call-X".equals(rr.id()))); + assertEquals(hasCall, hasResp, + "tool_call and matching response must both survive or both be dropped — " + + "orphan removal + pull-back guarantee this post-condition " + + "(hasCall=" + hasCall + ", hasResp=" + hasResp + ")"); + } + + // ---- absolute max safety net ------------------------------------------ + + @Test + @DisplayName("absoluteMax: pathological count is capped even when token budget allowed more") + void absoluteMax_caps() { + // 500 tiny messages — each ~5 chars; total tokens well under the + // budget so the token cut allows everything, but message count is + // pathological and must be capped. Anchor is at index 1 (right after + // the system header), so the cap would drop it without the stitch. + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("ANCHOR")); + for (int i = 0; i < 500; i++) msgs.add(new AssistantMessage("x")); + + // Larger minTailMessages floor not relevant here — we want to verify + // the hard count cap, which dominates over the floor in this case. + LoopBudgetConfig cfg = new LoopBudgetConfig(50_000, 30_000, 4, 1.5, 0, 50); + + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + assertTrue(r.trace().modified()); + assertTrue(r.trace().finalCount() <= 50, + "target max should cap the count regardless of token budget (got " + + r.trace().finalCount() + ")"); + assertTrue(r.trace().targetMaxTripped(), "trace should record the trip"); + assertTrue(containsExact(r.messages(), "ANCHOR"), + "even when absoluteMax forces a hard drop, the latest UserMessage must " + + "remain — stitched in if necessary"); + } + + // ---- config validation -------------------------------------------------- + + @Test + @DisplayName("LoopBudgetConfig.forContext picks sensible ratios from a context window") + void forContext_ratios() { + LoopBudgetConfig c = LoopBudgetConfig.forContext(128_000); + assertEquals(64_000, c.triggerTokens()); + assertEquals(38_400, c.keepTailTokens()); + assertEquals(4, c.minTailMessages()); + assertEquals(1.5, c.tailSoftCeilingRatio(), 0.001); + assertEquals(0, c.reservedPrefixTokens()); + assertEquals(200, c.targetMaxMessages()); + } + + @Test + @DisplayName("LoopBudgetConfig rejects degenerate values") + void config_rejectsBadValues() { + assertThrows(IllegalArgumentException.class, + () -> new LoopBudgetConfig(5_000, 6_000, 4, 1.5, 0, 200), + "tail must be strictly less than trigger"); + assertThrows(IllegalArgumentException.class, + () -> new LoopBudgetConfig(100, 50, 4, 1.5, 0, 200), + "trigger below MIN_TRIGGER_TOKENS rejected"); + assertThrows(IllegalArgumentException.class, + () -> new LoopBudgetConfig(50_000, 10_000, 1, 1.5, 0, 200), + "minTailMessages below floor rejected"); + assertThrows(IllegalArgumentException.class, + () -> new LoopBudgetConfig(50_000, 10_000, 4, 0.5, 0, 200), + "tailSoftCeilingRatio below 1.0 rejected"); + assertThrows(IllegalArgumentException.class, + () -> new LoopBudgetConfig(50_000, 10_000, 4, 1.5, -1, 200), + "negative reservedPrefixTokens rejected"); + } + + // ---- prefix-token accounting ------------------------------------------- + + @Test + @DisplayName("reservedPrefixTokens: trigger fires earlier when prefix is heavy") + void reservedPrefix_triggersBudgetWithSmallerHistory() { + // Build a moderate history that would NOT trigger on its own (~6K + // tokens), but combined with a 60K prefix reservation pushes the + // total over the 50K trigger. + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("anchor")); + for (int i = 0; i < 10; i++) msgs.add(new AssistantMessage(fat("obs-" + i))); + + // Without prefix accounting: well under 50K trigger. + LoopBudgetConfig noPrefix = new LoopBudgetConfig(50_000, 10_000, 4, 1.5, 0, 500); + assertFalse(budgeter.budget(msgs, noPrefix).trace().triggered(), + "without prefix reservation, this history must not trigger"); + + // With 60K prefix: should fire. + LoopBudgetConfig withPrefix = new LoopBudgetConfig(50_000, 10_000, 4, 1.5, 60_000, 500); + LoopMessageBudgeter.Result r = budgeter.budget(msgs, withPrefix); + assertTrue(r.trace().triggered(), + "with 60K reserved prefix, the same history must trip the trigger"); + assertEquals(60_000, r.trace().reservedPrefixTokens()); + } + + // ---- min-tail floor ---------------------------------------------------- + + @Test + @DisplayName("minTailMessages: a single huge tool output doesn't collapse the tail") + void minTail_floorPreservesRecentContext() { + // 50 fat history messages + an extremely fat last tool output that + // alone exceeds the hard tail budget. Without the floor, the + // backward walk in findTailCutByTokens would stop at just that one + // message, losing all recent reasoning. + // Anchor at the END so anchor enforcement doesn't paper over the + // floor — anchor-near-head would pull the cut back to include + // everything, masking the floor's structural effect. + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + for (int i = 0; i < 50; i++) msgs.add(new AssistantMessage(fat("head-fill-" + i))); + // The last assistant dwarfs the hard 5K tail budget. + msgs.add(new AssistantMessage(huge("monster-last"))); + msgs.add(new UserMessage("tail anchor")); + + LoopBudgetConfig cfg = new LoopBudgetConfig(30_000, 5_000, 4, 1.5, 0, 200); + + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + assertTrue(r.trace().triggered(), "token total exceeds trigger"); + assertTrue(r.trace().minTailFloorApplied(), + "the floor must have engaged — last message alone exceeds hard tail budget"); + assertTrue(r.trace().tailKept() >= 4, + "at least minTailMessages (4) entries should survive — got " + r.trace().tailKept()); + assertTrue(containsExact(r.messages(), "tail anchor"), + "anchor at the end must remain"); + } + + // ---- pair-integrity-vs-cap reporting ----------------------------------- + + @Test + @DisplayName("capExceededForPairIntegrity: cap is exceeded when pair pull-back wins") + void capExceeded_whenPairIntegrityForcesEarlierCut() { + // Setup: configure a targetMaxMessages cap such that the natural + // cap-driven cut lands inside a multi-response pair. Pair pull-back + // has to drag the boundary back across the entire pair → final + // count exceeds the cap, and the trace records the trade-off. + // Layout (idx): 0=sys, 1=anchor, 2-87=head fillers (86), + // 88=asst(c1..c5), 89-93=resp(c1..c5), 94-110=tail fillers (17). + // Total 111 messages. + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("anchor")); + for (int i = 0; i < 86; i++) msgs.add(new AssistantMessage(fat("head-" + i))); + msgs.add(asstMulti("c1", "c2", "c3", "c4", "c5")); + for (int i = 1; i <= 5; i++) msgs.add(resp("c" + i)); + for (int i = 0; i < 17; i++) msgs.add(new AssistantMessage("tail-" + i)); + + // Cap of 20 → tail cap = 19. provisionalTailStart = 111-19 = 92 + // (lands inside the responses). Pair pull-back drags it to 88. + // Final tail = 111-88 = 23, +head 1 = 24 > 20 cap → integrity wins. + LoopBudgetConfig cfg = new LoopBudgetConfig(20_000, 5_000, 4, 1.5, 0, 20); + + LoopMessageBudgeter.Result r = budgeter.budget(msgs, cfg); + + assertTrue(r.trace().targetMaxTripped()); + assertTrue(r.trace().capExceededForPairIntegrity(), + "pair pull-back forced final count above targetMaxMessages — " + + "trace must surface this so observers can react"); + assertTrue(r.trace().finalCount() > cfg.targetMaxMessages(), + "final count must actually exceed the cap (got " + r.trace().finalCount() + + ", cap=" + cfg.targetMaxMessages() + ")"); + // Pair integrity still held: every response has its call. + assertTrue(ToolPairSanitizer.isPaired(r.messages()), + "tool-pair invariant must hold post-budget"); + } + + // ---- ToolPairSanitizer post-condition ---------------------------------- + + @Test + @DisplayName("ToolPairSanitizer.isPaired: post-condition holds across all trim paths") + void sanitizer_postConditionAlwaysHolds() { + // Run several scenarios and assert the sanitizer post-condition. + // Building a quick matrix is cheaper than convincing ourselves the + // budgeter never produces an orphan, ever. + List> scenarios = new ArrayList<>(); + + // (a) clean history, no trim needed + List a = new ArrayList<>(); + a.add(new SystemMessage("sys")); + a.add(new UserMessage("u")); + a.add(asst("c1")); + a.add(resp("c1")); + a.add(new AssistantMessage("done")); + scenarios.add(a); + + // (b) trim with pair at boundary + List b = new ArrayList<>(); + b.add(new SystemMessage("sys")); + b.add(new UserMessage("u")); + for (int i = 0; i < 60; i++) b.add(new AssistantMessage(fat("h" + i))); + b.add(asst("c1")); + b.add(resp("c1")); + b.add(new AssistantMessage("tail")); + scenarios.add(b); + + // (c) cap-driven cut deep inside pair territory + List c = new ArrayList<>(); + c.add(new SystemMessage("sys")); + c.add(new UserMessage("u")); + for (int i = 0; i < 100; i++) c.add(new AssistantMessage("tiny" + i)); + c.add(asst("cx")); + for (int i = 0; i < 100; i++) c.add(new AssistantMessage("tiny-mid" + i)); + c.add(resp("cx")); + scenarios.add(c); + + for (int idx = 0; idx < scenarios.size(); idx++) { + LoopBudgetConfig cfg = new LoopBudgetConfig(8_000, 2_500, 4, 1.5, 0, 30); + LoopMessageBudgeter.Result r = budgeter.budget(scenarios.get(idx), cfg); + assertTrue(ToolPairSanitizer.isPaired(r.messages()), + "scenario " + idx + " produced an unpaired list: " + r.trace()); + } + } + + // ---- helpers ----------------------------------------------------------- + + private static LoopBudgetConfig defaultCfg() { + return LoopBudgetConfig.forContext(128_000); + } + + /** Produce a ~2KB string so token-budget tests cross thresholds with few entries. */ + private static String fat(String tag) { + StringBuilder sb = new StringBuilder(2000); + sb.append(tag).append(": "); + while (sb.length() < 2000) sb.append("lorem ipsum dolor sit amet "); + return sb.toString(); + } + + /** Produce a ~30KB string — bigger than typical tail budgets so it forces the floor. */ + private static String huge(String tag) { + StringBuilder sb = new StringBuilder(30_000); + sb.append(tag).append(": "); + while (sb.length() < 30_000) sb.append("lorem ipsum dolor sit amet consectetur "); + return sb.toString(); + } + + private static AssistantMessage asst(String callId) { + return AssistantMessage.builder().content("") + .toolCalls(List.of( + new AssistantMessage.ToolCall(callId, "function", "tool_" + callId, "{}"))) + .build(); + } + + private static AssistantMessage asstMulti(String... callIds) { + List calls = new ArrayList<>(); + for (String id : callIds) { + calls.add(new AssistantMessage.ToolCall(id, "function", "tool_" + id, "{}")); + } + return AssistantMessage.builder().content("").toolCalls(calls).build(); + } + + private static ToolResponseMessage resp(String callId) { + return ToolResponseMessage.builder().responses(List.of( + new ToolResponseMessage.ToolResponse(callId, "tool_" + callId, "ok"))).build(); + } + + /** Returns true if any UserMessage in {@code list} has text containing the given substring. */ + private static boolean containsExact(List list, String substring) { + return list.stream().anyMatch(m -> m instanceof UserMessage u && u.getText().contains(substring)); + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/agent/context/ToolPairSanitizerTest.java b/mateclaw-server/src/test/java/vip/mate/agent/context/ToolPairSanitizerTest.java new file mode 100644 index 00000000..99a46598 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/context/ToolPairSanitizerTest.java @@ -0,0 +1,231 @@ +package vip.mate.agent.context; + +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; +import org.springframework.ai.chat.messages.UserMessage; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * {@link ToolPairSanitizer} is the single source of truth for the + * tool_call ↔ tool_response pairing invariant. These tests exercise the + * three public methods directly against constructed message lists — the + * sanitizer is a pure-function utility, so there is no Spring context + * involved. + */ +class ToolPairSanitizerTest { + + // ---- isPaired() -------------------------------------------------------- + + @Test + @DisplayName("isPaired: empty list is trivially paired") + void isPaired_empty() { + assertTrue(ToolPairSanitizer.isPaired(List.of())); + assertTrue(ToolPairSanitizer.isPaired(null)); + } + + @Test + @DisplayName("isPaired: matched call and response are paired") + void isPaired_matched() { + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new UserMessage("u")); + msgs.add(asst("c1")); + msgs.add(resp("c1")); + assertTrue(ToolPairSanitizer.isPaired(msgs)); + } + + @Test + @DisplayName("isPaired: orphan response detected") + void isPaired_orphanResponse() { + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(resp("c1")); // no preceding call + assertFalse(ToolPairSanitizer.isPaired(msgs)); + } + + @Test + @DisplayName("isPaired: orphan call detected") + void isPaired_orphanCall() { + List msgs = new ArrayList<>(); + msgs.add(asst("c1")); // no following response + assertFalse(ToolPairSanitizer.isPaired(msgs)); + } + + @Test + @DisplayName("isPaired: null/empty id rejected as unpaired") + void isPaired_nullId() { + List msgs = new ArrayList<>(); + msgs.add(AssistantMessage.builder().content("") + .toolCalls(List.of(new AssistantMessage.ToolCall(null, "function", "t", "{}"))) + .build()); + assertFalse(ToolPairSanitizer.isPaired(msgs), + "an assistant tool_call with a null id can never be paired"); + + msgs.clear(); + msgs.add(ToolResponseMessage.builder().responses(List.of( + new ToolResponseMessage.ToolResponse("", "tool_x", "ok"))).build()); + assertFalse(ToolPairSanitizer.isPaired(msgs), + "a tool_response with an empty id can never be paired"); + } + + // ---- pullBackToToolPairBoundary() -------------------------------------- + + @Test + @DisplayName("pullBack: boundary inside a pair is moved before the assistant") + void pullBack_pairAtBoundary() { + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); // 0 + msgs.add(new AssistantMessage("pre")); // 1 + msgs.add(asst("c1")); // 2 — assistant tool_call + msgs.add(resp("c1")); // 3 — its response + msgs.add(new AssistantMessage("post")); // 4 + + // Proposed boundary at idx 3 would drop the assistant (idx 2) and + // keep the response (idx 3) → orphan. Pull-back should move the + // boundary back to idx 2 so the pair survives whole. + int adjusted = ToolPairSanitizer.pullBackToToolPairBoundary(msgs, 1, 3); + assertEquals(2, adjusted); + } + + @Test + @DisplayName("pullBack: boundary clear of any pair stays put") + void pullBack_noPairOverlap() { + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(new AssistantMessage("a")); + msgs.add(new AssistantMessage("b")); + msgs.add(new AssistantMessage("c")); + int adjusted = ToolPairSanitizer.pullBackToToolPairBoundary(msgs, 1, 2); + assertEquals(2, adjusted, "no tool pairs → boundary unchanged"); + } + + @Test + @DisplayName("pullBack: multi-call assistant with split responses pulls back across all of them") + void pullBack_multiCallAssistant() { + List msgs = new ArrayList<>(); + msgs.add(new SystemMessage("sys")); + msgs.add(asstMulti("c1", "c2", "c3")); // idx 1 — three calls + msgs.add(resp("c1")); // idx 2 + msgs.add(resp("c2")); // idx 3 + msgs.add(resp("c3")); // idx 4 + + // Boundary at idx 3 would keep c2 and c3 responses, orphaning them + // because their assistant (idx 1) would be dropped. + int adjusted = ToolPairSanitizer.pullBackToToolPairBoundary(msgs, 0, 3); + assertEquals(1, adjusted, + "boundary pulled to idx 1 so the multi-call assistant + all its responses survive"); + } + + // ---- removeOrphans() --------------------------------------------------- + + @Test + @DisplayName("removeOrphans: orphan response with no matching call is removed") + void removeOrphans_orphanResponse() { + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(resp("c1")); // orphan — no preceding call + + int removed = ToolPairSanitizer.removeOrphans(msgs); + + assertEquals(1, removed); + assertEquals(1, msgs.size()); + assertTrue(ToolPairSanitizer.isPaired(msgs)); + } + + @Test + @DisplayName("removeOrphans: assistant with NO matching responses removed") + void removeOrphans_orphanAssistant() { + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(asst("c1")); // orphan — no following response + + int removed = ToolPairSanitizer.removeOrphans(msgs); + + assertEquals(1, removed); + assertTrue(ToolPairSanitizer.isPaired(msgs)); + } + + @Test + @DisplayName("removeOrphans: assistant with PARTIAL matches is kept (lenient policy)") + void removeOrphans_partialMatchKept() { + // Lenient: assistant has c1 (matched) and c2 (orphan). Keep the + // whole assistant to preserve the matched pair; let strict provider + // dedup-handle the extra call rather than risking dropping useful + // history. + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(asstMulti("c1", "c2")); + msgs.add(resp("c1")); + + int removed = ToolPairSanitizer.removeOrphans(msgs); + + assertEquals(0, removed, + "an assistant with at least one matched call survives — partial-match lenient policy"); + assertEquals(3, msgs.size()); + } + + @Test + @DisplayName("removeOrphans: iterative — removing P1 reveals P0, both cleared") + void removeOrphans_iterativeConvergence() { + // Build a list where removing an orphan assistant (P1) leaves a now- + // dangling response (P0) that the next pass must also remove. + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(asst("c1")); // P1 — no response + msgs.add(resp("c2")); // P0 — no call (independent of c1) + + int removed = ToolPairSanitizer.removeOrphans(msgs); + + assertEquals(2, removed); + assertTrue(ToolPairSanitizer.isPaired(msgs)); + assertEquals(1, msgs.size()); + } + + @Test + @DisplayName("removeOrphans: matched pair untouched") + void removeOrphans_noOpOnCleanList() { + List msgs = new ArrayList<>(); + msgs.add(new UserMessage("u")); + msgs.add(asst("c1")); + msgs.add(resp("c1")); + msgs.add(new AssistantMessage("final")); + + int removed = ToolPairSanitizer.removeOrphans(msgs); + + assertEquals(0, removed); + assertEquals(4, msgs.size()); + assertTrue(ToolPairSanitizer.isPaired(msgs)); + } + + // ---- helpers ----------------------------------------------------------- + + private static AssistantMessage asst(String callId) { + return AssistantMessage.builder().content("") + .toolCalls(List.of( + new AssistantMessage.ToolCall(callId, "function", "tool_" + callId, "{}"))) + .build(); + } + + private static AssistantMessage asstMulti(String... callIds) { + List calls = new ArrayList<>(); + for (String id : callIds) { + calls.add(new AssistantMessage.ToolCall(id, "function", "tool_" + id, "{}")); + } + return AssistantMessage.builder().content("").toolCalls(calls).build(); + } + + private static ToolResponseMessage resp(String callId) { + return ToolResponseMessage.builder().responses(List.of( + new ToolResponseMessage.ToolResponse(callId, "tool_" + callId, "ok"))).build(); + } +}