From ddbd37b51b3b60087cf9b0a68438cd749b720fa1 Mon Sep 17 00:00:00 2001 From: matevip Date: Tue, 20 Oct 2026 04:03:19 -0400 Subject: [PATCH] fix: unify prompt budgets on AgentScope 2.0.4 --- .../context/BudgetedCompactionModel.java | 16 +---- .../vip/mate/agent/context/PromptBudget.java | 52 +++++++++++++++++ .../mate/agent/graph/node/ReasoningNode.java | 11 ++-- .../context/BudgetedCompactionModelTest.java | 26 +++++++++ .../mate/agent/context/PromptBudgetTest.java | 58 +++++++++++++++++++ 5 files changed, 143 insertions(+), 20 deletions(-) create mode 100644 mateclaw-server/src/main/java/vip/mate/agent/context/PromptBudget.java create mode 100644 mateclaw-server/src/test/java/vip/mate/agent/context/PromptBudgetTest.java diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/BudgetedCompactionModel.java b/mateclaw-server/src/main/java/vip/mate/agent/context/BudgetedCompactionModel.java index 4ba2c631..ab656f3d 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/context/BudgetedCompactionModel.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/BudgetedCompactionModel.java @@ -3,7 +3,6 @@ package vip.mate.agent.context; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; -import org.springframework.ai.model.tool.ToolCallingChatOptions; import vip.mate.agent.runtime.state.RunState; import vip.mate.llm.chatmodel.ChatModelDefaults; import java.util.concurrent.CancellationException; @@ -44,13 +43,7 @@ public final class BudgetedCompactionModel implements ChatModel { if (value != null && value>0) usage.merge(key,value,Math::addExact); } public static int outputReserveTokens(ChatOptions options,int fallback) { - Integer maximum = options == null ? null : options.getMaxTokens(); - if (options instanceof org.springframework.ai.openai.OpenAiChatOptions openAi - && openAi.getMaxCompletionTokens() != null) { - maximum = maximum == null ? openAi.getMaxCompletionTokens() - : Math.max(maximum,openAi.getMaxCompletionTokens()); - } - return maximum == null ? Math.max(0,fallback) : Math.max(0,maximum); + return PromptBudget.outputReserveTokens(options, null, fallback); } public BudgetedCompactionModel(ChatModel delegate,long allowedCalls,int windowTokens,int fallbackOutputTokens) { this.delegate=delegate; this.allowedCalls=allowedCalls; @@ -63,12 +56,7 @@ public final class BudgetedCompactionModel implements ChatModel { } private void checkPrompt(Prompt prompt) { if (windowTokens <= 0) return; - ChatOptions defaults = ChatModelDefaults.optionsFor(delegate); - ChatOptions options = prompt.getOptions() != null ? prompt.getOptions() : defaults; - int output = outputReserveTokens(options, outputReserveTokens(defaults,fallbackOutputTokens)); - int tools = options instanceof ToolCallingChatOptions toolOptions - ? TokenEstimator.estimateToolsTokens(toolOptions.getToolCallbacks()) : 0; - int total = TokenEstimator.estimateTokens(prompt.getInstructions()) + tools + output; + long total = PromptBudget.estimateTokens(prompt, ChatModelDefaults.optionsFor(delegate), fallbackOutputTokens); if (total > windowTokens) { blockedPrompt = "摘要 Prompt 估算 " + total + " tokens 超过上下文窗口 " + windowTokens + " tokens；未调用摘要模型。"; diff --git a/mateclaw-server/src/main/java/vip/mate/agent/context/PromptBudget.java b/mateclaw-server/src/main/java/vip/mate/agent/context/PromptBudget.java new file mode 100644 index 00000000..b340530f --- /dev/null +++ b/mateclaw-server/src/main/java/vip/mate/agent/context/PromptBudget.java @@ -0,0 +1,52 @@ +package vip.mate.agent.context; + +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.tool.ToolCallingChatOptions; +import org.springframework.ai.openai.OpenAiChatOptions; +import org.springframework.ai.tool.ToolCallback; + +import java.util.List; + +/** Estimates the effective request without changing the provider's prompt or options. */ +public final class PromptBudget { + private PromptBudget() {} + + public static long estimateTokens(Prompt prompt, ChatOptions defaults, int fallbackOutputTokens) { + ChatOptions runtime = prompt.getOptions(); + ChatOptions effective = effectiveOptions(runtime, defaults); + // Native combineWith treats empty lists as absent. Provider adapters explicitly replace + // callbacks when the runtime list is non-null, including a deliberately empty list. + List callbacks = runtime instanceof ToolCallingChatOptions tools + && tools.getToolCallbacks() != null ? tools.getToolCallbacks() + : effective instanceof ToolCallingChatOptions tools ? tools.getToolCallbacks() : List.of(); + return (long) TokenEstimator.estimateTokens(prompt.getInstructions()) + + TokenEstimator.estimateToolsTokens(callbacks) + + effectiveOutputReserve(effective, runtime, fallbackOutputTokens); + } + + public static int outputReserveTokens(ChatOptions runtime, ChatOptions defaults, int fallback) { + return effectiveOutputReserve(effectiveOptions(runtime, defaults), runtime, fallback); + } + + private static int effectiveOutputReserve(ChatOptions effective, ChatOptions runtime, int fallback) { + // A generic defaults builder cannot represent provider-specific runtime fields. Preserve + // the runtime's declared maximum as a conservative floor, as the existing guard did. + return Math.max(outputReserveTokens(effective, fallback), outputReserveTokens(runtime, 0)); + } + + private static ChatOptions effectiveOptions(ChatOptions runtime, ChatOptions defaults) { + if (defaults == null) return runtime; + if (runtime == null) return defaults; + return defaults.mutate().combineWith(runtime.mutate()).build(); + } + + private static int outputReserveTokens(ChatOptions options, int fallback) { + Integer maximum = options == null ? null : options.getMaxTokens(); + if (options instanceof OpenAiChatOptions openAi && openAi.getMaxCompletionTokens() != null) { + maximum = maximum == null ? openAi.getMaxCompletionTokens() + : Math.max(maximum, openAi.getMaxCompletionTokens()); + } + return maximum == null ? Math.max(0, fallback) : Math.max(0, maximum); + } +} 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 74f25e99..decc7716 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 @@ -1023,8 +1023,7 @@ public class ReasoningNode implements NodeAction { // budgeting retains legacy behavior. boolean mandatoryPrefixOverflow = loopReservedPrefixTokens + protectedHistoryTokens > loopContextWindowTokens(); - int assembledPromptEstimateTokens = TokenEstimator.estimateTokens(promptMessages) - + toolSchemaEstimateTokens + outputReserveTokens; + long assembledPromptEstimateTokens = vip.mate.agent.context.PromptBudget.estimateTokens(prompt, null, outputReserveTokens); boolean promptOverflow = (holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled()) && assembledPromptEstimateTokens > loopContextWindowTokens(); boolean localPromptRejected = promptOverflow; @@ -1149,8 +1148,8 @@ public class ReasoningNode implements NodeAction { // Reuse the SAME non-history prefix — wiki/runtime context preserved. List retryPromptMessages = new ArrayList<>(nonHistoryPrefix); retryPromptMessages.addAll(compactedMessages); - int retryEstimateTokens = TokenEstimator.estimateTokens(retryPromptMessages) - + toolSchemaEstimateTokens + outputReserveTokens; + long retryEstimateTokens = vip.mate.agent.context.PromptBudget.estimateTokens( + new Prompt(retryPromptMessages, options), null, outputReserveTokens); if ((holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled()) && retryEstimateTokens > loopContextWindowTokens()) { localPromptRejected = true; @@ -1234,8 +1233,8 @@ public class ReasoningNode implements NodeAction { List nudgeHistory = LOOP_BUDGETER.budget(messages, nudgeConfig).messages(); List nudgedMessages = new ArrayList<>(nudgePrefix); nudgedMessages.addAll(nudgeHistory); - int nudgeEstimate = TokenEstimator.estimateTokens(nudgedMessages) - + toolSchemaEstimateTokens + outputReserveTokens; + long nudgeEstimate = vip.mate.agent.context.PromptBudget.estimateTokens( + new Prompt(nudgedMessages, options), null, outputReserveTokens); if ((holdContent || prefixBudgetPlan != null && prefixBudgetPlan.enabled()) && nudgeEstimate > loopContextWindowTokens()) { localPromptRejected = true; diff --git a/mateclaw-server/src/test/java/vip/mate/agent/context/BudgetedCompactionModelTest.java b/mateclaw-server/src/test/java/vip/mate/agent/context/BudgetedCompactionModelTest.java index 2c5d57ab..df9820ea 100644 --- a/mateclaw-server/src/test/java/vip/mate/agent/context/BudgetedCompactionModelTest.java +++ b/mateclaw-server/src/test/java/vip/mate/agent/context/BudgetedCompactionModelTest.java @@ -12,6 +12,32 @@ import static org.mockito.Mockito.*; import static org.mockito.ArgumentMatchers.*; import static vip.mate.agent.graph.state.MateClawStateKeys.*; class BudgetedCompactionModelTest { + @Test void veryLargeOutputReserveCannotOverflowTotalAndBypassGuard() { + var delegate=mock(ChatModel.class); + var prompt=new Prompt("summary",ChatOptions.builder().maxTokens(Integer.MAX_VALUE).build()); + var guarded=new BudgetedCompactionModel(delegate,2,24000,1024); + assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.call(prompt)); + assertEquals(0,guarded.calls());verify(delegate,never()).call(any(Prompt.class)); + } + @Test void inheritedCompletionReserveIsCountedWhenRuntimeOverridesOnlyMaxTokens() { + var delegate=mock(ChatModel.class); + when(delegate.getDefaultOptions()).thenReturn(org.springframework.ai.openai.OpenAiChatOptions.builder().maxCompletionTokens(30000).build()); + var prompt=new Prompt("summary",org.springframework.ai.openai.OpenAiChatOptions.builder().maxTokens(512).build()); + var guarded=new BudgetedCompactionModel(delegate,2,24000,1024); + assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.call(prompt)); + assertEquals(0,guarded.calls());verify(delegate,never()).call(any(Prompt.class)); + } + @Test void inheritedDefaultToolSchemasAreCountedBeforeStreamingProviderInvocation() { + var delegate=mock(ChatModel.class);var tool=mock(org.springframework.ai.tool.ToolCallback.class); + when(delegate.stream(any(Prompt.class))).thenReturn(reactor.core.publisher.Flux.empty()); + when(tool.getToolDefinition()).thenReturn(org.springframework.ai.tool.definition.ToolDefinition.builder() + .name("large").description("界".repeat(1000)).inputSchema("{\"type\":\"object\"}").build()); + when(delegate.getDefaultOptions()).thenReturn(org.springframework.ai.openai.OpenAiChatOptions.builder().toolCallbacks(List.of(tool)).maxTokens(64).build()); + var prompt=new Prompt("summary",org.springframework.ai.openai.OpenAiChatOptions.builder().maxTokens(64).build()); + var guarded=new BudgetedCompactionModel(delegate,2,1000,64); + assertThrows(BudgetedCompactionModel.SummaryPromptOverflow.class,()->guarded.stream(prompt).blockLast()); + assertEquals(0,guarded.calls());verify(delegate,never()).stream(any(Prompt.class)); + } @Test void exactInvocationAndAvailableUsageAccumulateWithoutMutatingInputState() { var delegate=mock(ChatModel.class); when(delegate.call(any(Prompt.class))).thenReturn(new ChatResponse(List.of(new Generation(new AssistantMessage("summary"))),ChatResponseMetadata.builder().usage(new DefaultUsage(19,11,30,null,7L,3L)).build())); diff --git a/mateclaw-server/src/test/java/vip/mate/agent/context/PromptBudgetTest.java b/mateclaw-server/src/test/java/vip/mate/agent/context/PromptBudgetTest.java new file mode 100644 index 00000000..bf609708 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/context/PromptBudgetTest.java @@ -0,0 +1,58 @@ +package vip.mate.agent.context; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.openai.OpenAiChatOptions; +import org.springframework.ai.tool.ToolCallback; +import org.springframework.ai.tool.definition.ToolDefinition; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +class PromptBudgetTest { + @Test void explicitEmptyCallbacksAndLowerOutputOverrideDefaultsWithoutMutatingOptions() { + var tool=mock(ToolCallback.class); + var defaults=OpenAiChatOptions.builder().maxTokens(4096).toolCallbacks(List.of(tool)).build(); + var runtime=OpenAiChatOptions.builder().maxTokens(64).toolCallbacks(List.of()).build(); + var prompt=new Prompt("summary",runtime); + + assertEquals(TokenEstimator.estimateTokens(prompt.getInstructions())+64, + PromptBudget.estimateTokens(prompt,defaults,1024)); + assertEquals(4096,defaults.getMaxTokens()); + assertEquals(List.of(tool),defaults.getToolCallbacks()); + assertEquals(64,runtime.getMaxTokens()); + assertTrue(runtime.getToolCallbacks().isEmpty()); + assertSame(runtime,prompt.getOptions()); + verifyNoInteractions(tool); + } + + @Test void inheritedCompletionFieldAndRuntimeOverrideFollowNativeOptionMerge() { + var defaults=OpenAiChatOptions.builder().maxCompletionTokens(30000).build(); + var partial=OpenAiChatOptions.builder().maxTokens(512).build(); + assertEquals(30000,PromptBudget.outputReserveTokens(partial,defaults,1024)); + var explicit=partial.mutate().maxCompletionTokens(128).build(); + assertNull(explicit.getMaxTokens(),"Native builder clears maxTokens when setting maxCompletionTokens"); + assertEquals(128,PromptBudget.outputReserveTokens(explicit,defaults,1024)); + assertEquals(30000,defaults.getMaxCompletionTokens()); + assertNull(partial.getMaxCompletionTokens()); + } + + @Test void providerSpecificRuntimeOutputFieldSurvivesGenericDefaultBuilder() { + var defaults=ChatOptions.builder().maxTokens(512).build(); + var runtime=OpenAiChatOptions.builder().maxCompletionTokens(30000).build(); + assertEquals(30000,PromptBudget.outputReserveTokens(runtime,defaults,1024)); + assertTrue(PromptBudget.estimateTokens(new Prompt("small",runtime),defaults,1024)>30000); + } + + @Test void nullAndGenericRuntimeOptionsInheritDefaultToolSchemas() { + var tool=mock(ToolCallback.class); + when(tool.getToolDefinition()).thenReturn(ToolDefinition.builder() + .name("tool").description("界".repeat(1000)).inputSchema("{\"type\":\"object\"}").build()); + var defaults=OpenAiChatOptions.builder().maxTokens(64).toolCallbacks(List.of(tool)).build(); + assertTrue(PromptBudget.estimateTokens(new Prompt("small"),defaults,64)>1000); + assertTrue(PromptBudget.estimateTokens(new Prompt("small",ChatOptions.builder().maxTokens(64).build()),defaults,64)>1000); + } +}