mirror of
https://gitee.com/mateos/mateclaw.git
synced 2026-09-13 03:13:41 +08:00
fix(agent): replace per-loop head/tail trim with anchored token-budget budgeter
This commit is contained in:
matevip
parent
6c8c490bd3
commit
7f5652b2f0
@ -170,6 +170,15 @@ public class ConversationWindowManager {
|
||||
* conversation in a compaction storm. */
|
||||
private final ConcurrentHashMap<String, Long> 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
|
||||
|
||||
@ -0,0 +1,140 @@
|
||||
package vip.mate.agent.context;
|
||||
|
||||
/**
|
||||
* Configuration for per-reasoning-loop message budgeting.
|
||||
*
|
||||
* <p>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.
|
||||
*
|
||||
* <p>Field semantics:
|
||||
* <ul>
|
||||
* <li>{@code triggerTokens} — token threshold above which budgeting kicks
|
||||
* in. Compared against {@code historyTokens + reservedPrefixTokens}
|
||||
* so the budgeter accounts for the full prompt the LLM will see,
|
||||
* not just the message list.</li>
|
||||
* <li>{@code keepTailTokens} — token budget reserved for the tail (recent
|
||||
* observations + the current user message). Scales with the model
|
||||
* window instead of relying on a fixed count.</li>
|
||||
* <li>{@code minTailMessages} — floor on the kept-tail count. Prevents a
|
||||
* single huge tool output from collapsing the tail to one message and
|
||||
* losing recent reasoning context.</li>
|
||||
* <li>{@code tailSoftCeilingRatio} — multiplier applied to
|
||||
* {@code keepTailTokens} when honoring the floor or pulling back to
|
||||
* keep a tool pair whole. Lets the tail overshoot the hard budget by
|
||||
* up to this factor before more aggressive cuts kick in.</li>
|
||||
* <li>{@code reservedPrefixTokens} — estimated tokens consumed by the
|
||||
* non-history portion of the prompt (system prompt, skill catalog,
|
||||
* runtime context, wiki injection, tool schemas, output reserve).
|
||||
* Surfaces these from the caller so the budget covers the whole
|
||||
* prompt, not just the message list.</li>
|
||||
* <li>{@code targetMaxMessages} — soft ceiling on the count fed to the
|
||||
* LLM. Best-effort: the budgeter may exceed it slightly to keep a
|
||||
* tool pair whole rather than orphan a call/response — that case is
|
||||
* reported via {@code BudgetTrace.capExceededForPairIntegrity}.</li>
|
||||
* </ul>
|
||||
*/
|
||||
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.
|
||||
*
|
||||
* <ul>
|
||||
* <li>trigger = 50% of the window — same threshold the multi-turn
|
||||
* compressor uses, so the two layers stay calibrated.</li>
|
||||
* <li>tail budget = 30% of the window.</li>
|
||||
* <li>minTailMessages = 4 — at least one full reasoning/action cycle
|
||||
* stays visible to the LLM no matter how big a single tool output is.</li>
|
||||
* <li>tailSoftCeilingRatio = 1.5 — let the tail overshoot by 50% when
|
||||
* enforcing the floor or pulling back to keep a tool pair whole.</li>
|
||||
* <li>reservedPrefixTokens = 0 — caller should override with the real
|
||||
* prefix estimate; left at 0 the budget still works but errs on
|
||||
* the side of triggering later than it should.</li>
|
||||
* <li>targetMaxMessages = 200 — well above a normal ReAct loop's 20–40
|
||||
* working messages, low enough to be a meaningful guard rail.</li>
|
||||
* </ul>
|
||||
*/
|
||||
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);
|
||||
}
|
||||
}
|
||||
@ -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:
|
||||
*
|
||||
* <ol>
|
||||
* <li><b>System prompt(s)</b> — all consecutive {@link SystemMessage}s at
|
||||
* the head stay verbatim. Production agents commonly have multiple
|
||||
* (SOUL, AGENTS, runtime context, wiki, tool prompt, skill catalog).</li>
|
||||
* <li><b>Turn anchor</b> — the latest {@link UserMessage} is never dropped.
|
||||
* Stitched in when an aggressive cut would otherwise lose it.</li>
|
||||
* <li><b>Tool-call/response pair integrity</b> — every assistant tool_call
|
||||
* reaches the model with its matching tool_response, and vice versa.
|
||||
* Delegated to {@link ToolPairSanitizer}.</li>
|
||||
* <li><b>Token budget over message count</b> — tail sized by token estimate
|
||||
* so a small ReAct loop with fat observations and a large loop with
|
||||
* thin observations both fit one config.</li>
|
||||
* <li><b>Minimum tail messages</b> — 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.</li>
|
||||
* </ol>
|
||||
*
|
||||
* <p>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).
|
||||
*
|
||||
* <p>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.
|
||||
*
|
||||
* <p>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<Message> messages, BudgetTrace trace) {}
|
||||
|
||||
/**
|
||||
* Structured trace of a single budgeting decision. All counts and token
|
||||
* figures refer to {@link Message} entries.
|
||||
*
|
||||
* <ul>
|
||||
* <li>{@code anchorEnforced} — the tail cut was pulled earlier than
|
||||
* the token budget would have placed it because the latest
|
||||
* UserMessage would otherwise have been dropped.</li>
|
||||
* <li>{@code anchorStitched} — the latest UserMessage could not fit in
|
||||
* the tail even after pull-back (typically when the target cap
|
||||
* fired hard); it was inserted as a standalone slot between head
|
||||
* and tail.</li>
|
||||
* <li>{@code capExceededForPairIntegrity} — the final count exceeded
|
||||
* {@code targetMaxMessages} because pulling the cut back to keep
|
||||
* a tool pair whole won out over the soft cap. Useful signal that
|
||||
* upstream compaction should have run sooner.</li>
|
||||
* <li>{@code minTailFloorApplied} — the tail was enlarged past the
|
||||
* hard token budget (up to the soft ceiling) to honor
|
||||
* {@code minTailMessages}.</li>
|
||||
* <li>{@code triggered} — the budget entered its main path because the
|
||||
* trigger threshold was met. Says nothing about whether anything
|
||||
* was actually removed.</li>
|
||||
* <li>{@code modified} — the returned list differs from the input
|
||||
* (count changed or orphans removed). This is the only signal
|
||||
* callers should use to gate log output; a triggered-but-no-op
|
||||
* pass is normal and shouldn't spam logs.</li>
|
||||
* </ul>
|
||||
*/
|
||||
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<Message> 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<Message> 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<Message> 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<Message> 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<Message> messages, int headEnd) {
|
||||
for (int i = messages.size() - 1; i >= headEnd; i--) {
|
||||
if (messages.get(i) instanceof UserMessage) {
|
||||
return i;
|
||||
}
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
@ -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:
|
||||
*
|
||||
* <ul>
|
||||
* <li>Every {@code tool_call.id} on an {@link AssistantMessage} has a
|
||||
* matching {@code tool_response.id} on a {@link ToolResponseMessage}
|
||||
* <em>after</em> it in the list.</li>
|
||||
* <li>Every {@code tool_response.id} on a {@link ToolResponseMessage} has
|
||||
* a matching {@code tool_call.id} on an {@link AssistantMessage}
|
||||
* <em>before</em> it.</li>
|
||||
* <li>No empty/null ids on either side.</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>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.
|
||||
*
|
||||
* <p>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<Message> messages, int headEnd, int tailStart) {
|
||||
if (tailStart <= headEnd || messages == null || messages.isEmpty()) {
|
||||
return tailStart;
|
||||
}
|
||||
Set<String> 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:
|
||||
*
|
||||
* <p><b>P0</b>: a {@link ToolResponseMessage} whose response id has no
|
||||
* matching assistant tool_call in the list.
|
||||
*
|
||||
* <p><b>P1</b>: 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.)
|
||||
*
|
||||
* <p>Iterates until convergence: removing an assistant for P1 can expose
|
||||
* a P0 orphan that needs cleaning, and vice versa.
|
||||
*
|
||||
* <p>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<Message> messages) {
|
||||
if (messages == null || messages.isEmpty()) {
|
||||
return 0;
|
||||
}
|
||||
int totalRemoved = 0;
|
||||
boolean changed;
|
||||
do {
|
||||
Set<String> callIds = new HashSet<>();
|
||||
Set<String> 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<Message> messages) {
|
||||
if (messages == null || messages.isEmpty()) {
|
||||
return true;
|
||||
}
|
||||
Set<String> callIds = new HashSet<>();
|
||||
Set<String> 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;
|
||||
}
|
||||
}
|
||||
@ -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 <em>after</em>
|
||||
* 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<Message> 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<Message> 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<String> callIds = new java.util.HashSet<>();
|
||||
Set<String> 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, "");
|
||||
|
||||
@ -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.
|
||||
*
|
||||
* <p>The single load-bearing invariant is the <b>anchor</b>: 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<List<Message>> scenarios = new ArrayList<>();
|
||||
|
||||
// (a) clean history, no trim needed
|
||||
List<Message> 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<Message> 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<Message> 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<AssistantMessage.ToolCall> 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<Message> list, String substring) {
|
||||
return list.stream().anyMatch(m -> m instanceof UserMessage u && u.getText().contains(substring));
|
||||
}
|
||||
}
|
||||
@ -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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<Message> 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<AssistantMessage.ToolCall> 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();
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user