From 210f465bced4741a66bc2662315c40fd44bf0821 Mon Sep 17 00:00:00 2001 From: konlue <2362586639@qq.com> Date: Mon, 2 Nov 2026 20:01:06 +0800 Subject: [PATCH 1/4] feat(agent): gate the final model request on a context window budget (issue #137) --- .../ApproximateRequestTokenEstimator.java | 143 +++++++++++ .../ai/agent/budget/BudgetComponent.java | 48 ++++ .../ai/agent/budget/BudgetViolation.java | 66 +++++ .../budget/ContextBudgetInterceptor.java | 164 +++++++++++++ .../ai/agent/budget/ContextWindowBudget.java | 60 +++++ .../ContextWindowBudgetExceededException.java | 83 +++++++ .../ai/agent/budget/EstimationAccuracy.java | 44 ++++ .../ai/agent/budget/RequestCompactor.java | 43 ++++ .../ai/agent/budget/RequestTokenEstimate.java | 89 +++++++ .../agent/budget/RequestTokenEstimator.java | 40 ++++ .../budget/ToolCallPairingValidator.java | 152 ++++++++++++ .../budget/ContextBudgetInterceptorTest.java | 225 ++++++++++++++++++ .../LargeResultEvictionInterceptor.java | 40 +++- .../SummarizationRequestCompactor.java | 77 ++++++ .../SummarizationRequestCompactorTest.java | 119 +++++++++ .../LargeResultEvictionInterceptorTest.java | 55 +++++ 16 files changed, 1440 insertions(+), 8 deletions(-) create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetComponent.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetViolation.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudget.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudgetExceededException.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/EstimationAccuracy.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestCompactor.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimate.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimator.java create mode 100644 argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java create mode 100644 argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java create mode 100644 argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/hook/messages/SummarizationRequestCompactor.java create mode 100644 argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/hooks/messages/SummarizationRequestCompactorTest.java diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java new file mode 100644 index 000000000..83c30cc58 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java @@ -0,0 +1,143 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ContentPart; +import io.github.agentic.ai.api.chat.ToolCall; +import io.github.agentic.ai.api.chat.ToolDefinition; + +import java.util.HashSet; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +/** + * The default estimator: a character based approximation of the token consumption of + * a request. + * + *

+ * Unlike the message-only counters of the framework, the estimate covers the whole + * request which is finally dispatched: the system message, the conversation, the + * disclosed tool schemas and the reserved output. + * + *

+ * The tool schemas are counted for the bindings the request actually discloses. When + * the request carries an explicit tool selection, only the selected bindings are + * counted; the dynamic bindings of the request are always counted because they are + * disclosed next to the selected ones. + */ +public final class ApproximateRequestTokenEstimator implements RequestTokenEstimator { + + /** + * The average number of characters per token used by the approximation. + */ + public static final int DEFAULT_CHARS_PER_TOKEN = 4; + + private final int charsPerToken; + + /** + * Creates an estimator with the default character ratio. + */ + public ApproximateRequestTokenEstimator() { + this(DEFAULT_CHARS_PER_TOKEN); + } + + /** + * Creates an estimator with a custom character ratio. + * @param charsPerToken the average number of characters per token, must be positive + */ + public ApproximateRequestTokenEstimator(int charsPerToken) { + if (charsPerToken <= 0) { + throw new IllegalArgumentException("charsPerToken must be positive"); + } + this.charsPerToken = charsPerToken; + } + + @Override + public RequestTokenEstimate estimate(ModelRequest request, ContextWindowBudget budget) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(budget, "budget must not be null"); + return new RequestTokenEstimate(systemTokens(request), messageTokens(request), toolSchemaTokens(request), + budget.outputReserve(), EstimationAccuracy.APPROXIMATE); + } + + private int systemTokens(ModelRequest request) { + ChatMessage systemMessage = request.getSystemMessage(); + return systemMessage == null ? 0 : messageTokens(systemMessage); + } + + private int messageTokens(ModelRequest request) { + int total = 0; + for (ChatMessage message : request.getMessages()) { + total += messageTokens(message); + } + return total; + } + + private int messageTokens(ChatMessage message) { + int total = 0; + for (ContentPart part : message.content()) { + if (part.text() != null) { + total += tokens(part.text()); + } + else if (part.data() != null) { + total += tokens(part.data().length); + } + } + for (ToolCall call : message.toolCalls()) { + total += tokens(call.arguments()); + } + return total; + } + + private int toolSchemaTokens(ModelRequest request) { + Set selection = request.hasToolSelection() ? new HashSet<>(request.getTools()) : null; + int total = 0; + for (ToolDefinition definition : request.getToolRegistry().definitions()) { + if (selection == null || selection.contains(definition.name())) { + total += definitionTokens(definition); + } + } + for (ToolDefinition definition : request.getDynamicToolRegistry().definitions()) { + total += definitionTokens(definition); + } + for (Map.Entry description : request.getToolDescriptions().entrySet()) { + total += tokens(description.getKey()) + tokens(description.getValue()); + } + return total; + } + + private int definitionTokens(ToolDefinition definition) { + return tokens(definition.name()) + tokens(definition.description()) + tokens(definition.inputSchema()); + } + + private int tokens(String text) { + if (text == null || text.isEmpty()) { + return 0; + } + return tokens(text.length()); + } + + private int tokens(int length) { + if (length <= 0) { + return 0; + } + return (length + charsPerToken - 1) / charsPerToken; + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetComponent.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetComponent.java new file mode 100644 index 000000000..874960080 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetComponent.java @@ -0,0 +1,48 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +/** + * The part of a request which exceeded the budget. + */ +public enum BudgetComponent { + + /** + * The system message. + */ + SYSTEM, + + /** + * The conversation. + */ + MESSAGES, + + /** + * The disclosed tool schemas. + */ + TOOL_SCHEMA, + + /** + * The tokens reserved for the completion. + */ + OUTPUT_RESERVE, + + /** + * The request as a whole. + */ + TOTAL + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetViolation.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetViolation.java new file mode 100644 index 000000000..87c9b6e18 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/BudgetViolation.java @@ -0,0 +1,66 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import java.util.Objects; + +/** + * A budget which was exceeded. + * + *

+ * The violation is a value, not a message: a caller can react to the classified + * component instead of parsing an error string. + * + * @param component the part which exceeded the budget + * @param estimatedTokens the estimated tokens of that part + * @param limitTokens the tokens available for it + */ +public record BudgetViolation(BudgetComponent component, int estimatedTokens, int limitTokens) { + + /** + * Creates a budget violation. + * @param component the part which exceeded the budget, must not be null + * @param estimatedTokens the estimated tokens of that part, must not be negative + * @param limitTokens the tokens available for it, must not be negative + */ + public BudgetViolation { + Objects.requireNonNull(component, "component must not be null"); + if (estimatedTokens < 0 || limitTokens < 0) { + throw new IllegalArgumentException("token counts must not be negative"); + } + } + + /** + * Describes the violation of the total budget. + * @param estimate the estimate which exceeded the budget, must not be null + * @param budget the budget which was exceeded, must not be null + * @return the violation + */ + public static BudgetViolation ofTotal(RequestTokenEstimate estimate, ContextWindowBudget budget) { + Objects.requireNonNull(estimate, "estimate must not be null"); + Objects.requireNonNull(budget, "budget must not be null"); + return new BudgetViolation(BudgetComponent.TOTAL, estimate.totalTokens(), budget.inputTokenLimit()); + } + + /** + * The tokens which exceeded the budget. + * @return the excess in tokens + */ + public int excessTokens() { + return Math.max(0, estimatedTokens - limitTokens); + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java new file mode 100644 index 000000000..b3bb6d975 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java @@ -0,0 +1,164 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.agent.interceptor.ModelCallHandler; +import io.github.agentic.ai.agent.interceptor.ModelInterceptor; +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.interceptor.ModelResponse; + +import java.util.Objects; +import java.util.Optional; + +/** + * The gate which keeps a model request inside the context window of the model. + * + *

+ * The interceptor runs at the last point before the request is dispatched. It + * estimates the token consumption of the whole request, and when the request does not + * fit: + * + *

    + *
  1. the configured {@link RequestCompactor} is asked to shrink the request;
    2. + *
  2. the compacted request is estimated again;
    3. + *
  3. the tool call pairing of the compacted request is validated;
    4. + *
  4. the request is dispatched only when it fits and stays paired.
    5. + *
+ * + *

+ * A request which still does not fit is never dispatched. The interceptor fails with + * a {@link ContextWindowBudgetExceededException} instead, so an oversized request + * cannot reach the provider. + * + *

+ * The interceptor is opt-in and does not change the behaviour of an agent unless it + * is registered in the interceptor chain. + */ +public final class ContextBudgetInterceptor extends ModelInterceptor { + + private final ContextWindowBudget budget; + + private final RequestTokenEstimator estimator; + + private final RequestCompactor compactor; + + private ContextBudgetInterceptor(Builder builder) { + this.budget = builder.budget; + this.estimator = builder.estimator; + this.compactor = builder.compactor; + } + + /** + * Creates a builder for the interceptor. + * @return a new builder + */ + public static Builder builder() { + return new Builder(); + } + + @Override + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler handler) { + Objects.requireNonNull(request, "request must not be null"); + Objects.requireNonNull(handler, "handler must not be null"); + RequestTokenEstimate estimate = estimator.estimate(request, budget); + if (estimate.totalTokens() <= budget.inputTokenLimit()) { + return handler.call(request); + } + ModelRequest compacted = request; + if (compactor != null) { + compacted = Objects.requireNonNull(compactor.compact(request), "compactor must not return null"); + Optional violation = ToolCallPairingValidator + .validate(compacted.getMessages()); + if (violation.isPresent()) { + throw new ToolCallPairingValidator.PairingException(violation.get()); + } + } + RequestTokenEstimate compactedEstimate = estimator.estimate(compacted, budget); + if (compactedEstimate.totalTokens() > budget.inputTokenLimit()) { + throw ContextWindowBudgetExceededException.of(compactedEstimate, budget); + } + return handler.call(compacted); + } + + @Override + public String getName() { + return "ContextBudget"; + } + + /** + * The budget this interceptor enforces. + * @return the budget, never null + */ + public ContextWindowBudget budget() { + return budget; + } + + /** + * Builder for {@link ContextBudgetInterceptor}. + */ + public static final class Builder { + + private ContextWindowBudget budget; + + private RequestTokenEstimator estimator = new ApproximateRequestTokenEstimator(); + + private RequestCompactor compactor; + + private Builder() { + } + + /** + * Sets the budget of the target model. + * @param budget the budget, must not be null + * @return this builder + */ + public Builder budget(ContextWindowBudget budget) { + this.budget = budget; + return this; + } + + /** + * Sets the estimator which measures the request. + * @param estimator the estimator, must not be null + * @return this builder + */ + public Builder estimator(RequestTokenEstimator estimator) { + this.estimator = estimator; + return this; + } + + /** + * Sets the compactor which shrinks an oversized request. + * @param compactor the compactor, nullable to disable compression + * @return this builder + */ + public Builder compactor(RequestCompactor compactor) { + this.compactor = compactor; + return this; + } + + /** + * Builds the interceptor. + * @return the interceptor + */ + public ContextBudgetInterceptor build() { + Objects.requireNonNull(budget, "budget must not be null"); + Objects.requireNonNull(estimator, "estimator must not be null"); + return new ContextBudgetInterceptor(this); + } + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudget.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudget.java new file mode 100644 index 000000000..341dcb9fa --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudget.java @@ -0,0 +1,60 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +/** + * The context window budget of a model. + * + *

+ * A budget has two parts. The {@code inputTokenLimit} is the total context window of + * the model, and the {@code outputReserve} is the part of that window which is kept + * for the completion. Everything the request carries - the system message, the + * conversation, the tool schemas and the reserved output - has to fit into the + * window together. + * + *

+ * The budget describes the window of one model, not of a run. A run that talks to + * several models carries one budget per model. + */ +public record ContextWindowBudget(int inputTokenLimit, int outputReserve) { + + /** + * Creates a context window budget. + * @param inputTokenLimit the total context window of the model, must be positive + * @param outputReserve the tokens reserved for the completion, must be positive and + * smaller than the window + */ + public ContextWindowBudget { + if (inputTokenLimit <= 0) { + throw new IllegalArgumentException("inputTokenLimit must be positive"); + } + if (outputReserve <= 0) { + throw new IllegalArgumentException("outputReserve must be positive"); + } + if (outputReserve >= inputTokenLimit) { + throw new IllegalArgumentException("outputReserve must be smaller than inputTokenLimit"); + } + } + + /** + * The number of tokens which may be spent on the request itself. + * @return the usable input budget + */ + public int usableInputTokens() { + return inputTokenLimit - outputReserve; + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudgetExceededException.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudgetExceededException.java new file mode 100644 index 000000000..530052703 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextWindowBudgetExceededException.java @@ -0,0 +1,83 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import java.util.Objects; + +/** + * Thrown when a request does not fit into the context window of the model and could + * not be compressed into it. + * + *

+ * The failure is explicit. The request is not dispatched, and it is not silently + * truncated either: a caller which cannot fit the request has to decide how to + * continue. + */ +public class ContextWindowBudgetExceededException extends RuntimeException { + + private final BudgetViolation violation; + + private final RequestTokenEstimate estimate; + + /** + * Creates the exception. + * @param violation the classified violation, must not be null + * @param estimate the estimate which exceeded the budget, must not be null + */ + public ContextWindowBudgetExceededException(BudgetViolation violation, RequestTokenEstimate estimate) { + super("the request does not fit into the context window: " + violation.estimatedTokens() + " estimated tokens, " + + violation.limitTokens() + " available, dominant part is " + estimate.dominantComponent()); + this.violation = violation; + this.estimate = estimate; + } + + /** + * The classified violation. + * @return the violation, never null + */ + public BudgetViolation violation() { + return violation; + } + + /** + * The estimate which exceeded the budget. + * @return the estimate, never null + */ + public RequestTokenEstimate estimate() { + return estimate; + } + + /** + * The part of the request which consumes the most tokens. + * @return the dominant component, never null + */ + public BudgetComponent dominantComponent() { + return estimate.dominantComponent(); + } + + /** + * Creates the exception for an estimate which exceeded a budget. + * @param estimate the estimate which exceeded the budget, must not be null + * @param budget the budget which was exceeded, must not be null + * @return the exception + */ + public static ContextWindowBudgetExceededException of(RequestTokenEstimate estimate, ContextWindowBudget budget) { + Objects.requireNonNull(estimate, "estimate must not be null"); + Objects.requireNonNull(budget, "budget must not be null"); + return new ContextWindowBudgetExceededException(BudgetViolation.ofTotal(estimate, budget), estimate); + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/EstimationAccuracy.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/EstimationAccuracy.java new file mode 100644 index 000000000..a695be92f --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/EstimationAccuracy.java @@ -0,0 +1,44 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +/** + * How accurate a token estimate is. + * + *

+ * The distinction matters because a budget gate that believes it knows the exact + * token count behaves differently from a gate that only approximates it. A caller + * which cannot tolerate an approximation has to plug in an exact estimator instead + * of trusting a heuristic. + */ +public enum EstimationAccuracy { + + /** + * The count is the exact number of tokens the provider will consume. + */ + EXACT, + + /** + * The count is a heuristic and may be lower or higher than the real usage. + */ + APPROXIMATE, + + /** + * The count never underestimates the real usage. + */ + CONSERVATIVE + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestCompactor.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestCompactor.java new file mode 100644 index 000000000..d82e52c8a --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestCompactor.java @@ -0,0 +1,43 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.agent.interceptor.ModelRequest; + +/** + * Shrinks a request which does not fit into the context window of the model. + * + *

+ * The compactor is the seam between the budget and the existing compression + * capabilities. It does not define how a request is compressed; it only defines the + * result: a request which is not larger than the one it received. + * + *

+ * A compactor which cannot compress anything returns the request unchanged. The + * budget gate then re-estimates the request and fails explicitly instead of + * dispatching an oversized request. + */ +@FunctionalInterface +public interface RequestCompactor { + + /** + * Compresses a request. + * @param request the request to compress, never null + * @return the compressed request, never null + */ + ModelRequest compact(ModelRequest request); + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimate.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimate.java new file mode 100644 index 000000000..9b8b73a21 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimate.java @@ -0,0 +1,89 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import java.util.Objects; + +/** + * The estimated token consumption of one model request. + * + *

+ * The estimate is decomposed instead of being a single number, because the reaction + * to an oversized request depends on which part caused it: a large tool schema is + * handled by selecting fewer tools, a large conversation by summarizing it, and a + * large output reserve by asking for a shorter answer. + * + * @param systemTokens the estimated tokens of the system message + * @param messageTokens the estimated tokens of the conversation + * @param toolSchemaTokens the estimated tokens of the disclosed tool schemas + * @param outputReserve the tokens reserved for the completion + * @param accuracy how accurate the estimate is + */ +public record RequestTokenEstimate(int systemTokens, int messageTokens, int toolSchemaTokens, int outputReserve, + EstimationAccuracy accuracy) { + + /** + * Creates a token estimate. + * @param systemTokens the estimated tokens of the system message, must not be + * negative + * @param messageTokens the estimated tokens of the conversation, must not be + * negative + * @param toolSchemaTokens the estimated tokens of the disclosed tool schemas, must + * not be negative + * @param outputReserve the tokens reserved for the completion, must not be negative + * @param accuracy how accurate the estimate is, must not be null + */ + public RequestTokenEstimate { + if (systemTokens < 0 || messageTokens < 0 || toolSchemaTokens < 0 || outputReserve < 0) { + throw new IllegalArgumentException("token estimates must not be negative"); + } + Objects.requireNonNull(accuracy, "accuracy must not be null"); + } + + /** + * The estimated tokens of the request itself, without the reserved output. + * @return the estimated input tokens + */ + public int inputTokens() { + return systemTokens + messageTokens + toolSchemaTokens; + } + + /** + * The estimated tokens the request and the reserved output consume together. + * @return the estimated total tokens + */ + public int totalTokens() { + return inputTokens() + outputReserve; + } + + /** + * The part of the request which consumes the most tokens. + * @return the dominant component, never null + */ + public BudgetComponent dominantComponent() { + if (toolSchemaTokens >= messageTokens && toolSchemaTokens >= systemTokens && toolSchemaTokens >= outputReserve) { + return BudgetComponent.TOOL_SCHEMA; + } + if (messageTokens >= systemTokens && messageTokens >= outputReserve) { + return BudgetComponent.MESSAGES; + } + if (systemTokens >= outputReserve) { + return BudgetComponent.SYSTEM; + } + return BudgetComponent.OUTPUT_RESERVE; + } + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimator.java new file mode 100644 index 000000000..a15d61930 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/RequestTokenEstimator.java @@ -0,0 +1,40 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.agent.interceptor.ModelRequest; + +/** + * Estimates how many tokens a model request consumes. + * + *

+ * The estimator is the replaceable part of the budget contract. The default + * implementation approximates, a provider specific implementation can be exact. The + * budget gate does not care which one it is given, it only reads the resulting + * estimate and its accuracy. + */ +@FunctionalInterface +public interface RequestTokenEstimator { + + /** + * Estimates the token consumption of a request. + * @param request the request to estimate, must not be null + * @param budget the budget of the target model, must not be null + * @return the estimate, never null + */ + RequestTokenEstimate estimate(ModelRequest request, ContextWindowBudget budget); + +} diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java new file mode 100644 index 000000000..baab428e8 --- /dev/null +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java @@ -0,0 +1,152 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ChatRole; +import io.github.agentic.ai.api.chat.ToolCall; + +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +/** + * Validates that the tool calls of a conversation and their results stay paired. + * + *

+ * Compression may remove a message. When it removes a tool result, the assistant + * message which requested the tool call is left with an unanswered call; when it + * removes the assistant message, the tool result is left without a call. Both break + * the conversation, so the budget gate refuses to dispatch such a request. + */ +public final class ToolCallPairingValidator { + + private ToolCallPairingValidator() { + } + + /** + * The kind of a pairing violation. + */ + public enum Kind { + + /** + * A tool call has no matching tool result. + */ + MISSING_TOOL_RESULT, + + /** + * A tool result has no matching tool call. + */ + UNKNOWN_TOOL_RESULT, + + /** + * The same tool call appears more than once. + */ + DUPLICATE_TOOL_CALL + } + + /** + * A pairing violation. + * + * @param kind the kind of the violation, never null + * @param toolCallId the identifier of the affected tool call, never null + */ + public record Violation(Kind kind, String toolCallId) { + + /** + * Creates a pairing violation. + * @param kind the kind of the violation, must not be null + * @param toolCallId the identifier of the affected tool call, must not be null + */ + public Violation { + Objects.requireNonNull(kind, "kind must not be null"); + Objects.requireNonNull(toolCallId, "toolCallId must not be null"); + } + } + + /** + * Validates the pairing of the tool calls and the tool results of a + * conversation. + * @param messages the conversation, must not be null + * @return the violation, or empty when the conversation is paired + */ + public static Optional validate(List messages) { + Objects.requireNonNull(messages, "messages must not be null"); + Map callCounts = new HashMap<>(); + Set calls = new LinkedHashSet<>(); + for (ChatMessage message : messages) { + for (ToolCall call : message.toolCalls()) { + if (call.id() == null) { + continue; + } + callCounts.merge(call.id(), 1, Integer::sum); + calls.add(call.id()); + } + } + for (Map.Entry call : callCounts.entrySet()) { + if (call.getValue() > 1) { + return Optional.of(new Violation(Kind.DUPLICATE_TOOL_CALL, call.getKey())); + } + } + Set results = new LinkedHashSet<>(); + for (ChatMessage message : messages) { + if (message.role() == ChatRole.TOOL && message.toolCallId() != null) { + results.add(message.toolCallId()); + } + } + for (String result : results) { + if (!calls.contains(result)) { + return Optional.of(new Violation(Kind.UNKNOWN_TOOL_RESULT, result)); + } + } + for (String call : calls) { + if (!results.contains(call)) { + return Optional.of(new Violation(Kind.MISSING_TOOL_RESULT, call)); + } + } + return Optional.empty(); + } + + /** + * Thrown when a request which is about to be dispatched has a broken pairing. + */ + public static final class PairingException extends RuntimeException { + + private final Violation violation; + + /** + * Creates the exception. + * @param violation the violation, must not be null + */ + public PairingException(Violation violation) { + super("tool call pairing is broken: " + violation.kind() + " (" + violation.toolCallId() + ")"); + this.violation = violation; + } + + /** + * The classified violation. + * @return the violation, never null + */ + public Violation violation() { + return violation; + } + } + +} diff --git a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java new file mode 100644 index 000000000..4d0cca4b5 --- /dev/null +++ b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java @@ -0,0 +1,225 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.interceptor.ModelResponse; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ToolCall; +import io.github.agentic.ai.api.chat.ToolDefinition; +import io.github.agentic.ai.api.tool.Tool; +import io.github.agentic.ai.api.tool.ToolContext; +import io.github.agentic.ai.api.tool.ToolRegistry; +import io.github.agentic.ai.api.tool.ToolResult; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionStage; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Contract tests for the context window budget gate. + * + *

+ * The tests are model independent: no provider is called, the request is measured + * and the dispatch is counted. The invariant under test is that an oversized request + * never reaches the provider. + */ +class ContextBudgetInterceptorTest { + + @Test + void theEstimateCoversSystemMessagesToolSchemasAndTheOutputReserve() { + RequestTokenEstimate estimate = new ApproximateRequestTokenEstimator().estimate(request(), budget()); + + assertThat(estimate.systemTokens()).isPositive(); + assertThat(estimate.messageTokens()).isPositive(); + assertThat(estimate.toolSchemaTokens()).isPositive(); + assertThat(estimate.outputReserve()).isEqualTo(100); + assertThat(estimate.totalTokens()).isEqualTo(estimate.inputTokens() + 100); + assertThat(estimate.accuracy()).isEqualTo(EstimationAccuracy.APPROXIMATE); + } + + @Test + void anOversizedRequestIsNeverDispatched() { + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(20, 5)) + .build(); + + assertThatThrownBy(() -> interceptor.interceptModel(request(), r -> { + dispatches.incrementAndGet(); + return response(); + })).isInstanceOf(ContextWindowBudgetExceededException.class) + .satisfies(failure -> { + ContextWindowBudgetExceededException exceeded = (ContextWindowBudgetExceededException) failure; + assertThat(exceeded.violation().component()).isEqualTo(BudgetComponent.TOTAL); + assertThat(exceeded.violation().excessTokens()).isPositive(); + assertThat(exceeded.dominantComponent()).isNotNull(); + }); + + assertThat(dispatches).hasValue(0); + } + + @Test + void aCompactedRequestIsMeasuredAgainBeforeItIsDispatched() { + AtomicInteger dispatches = new AtomicInteger(); + AtomicReference dispatched = new AtomicReference<>(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(100, 5)) + .compactor(request -> ModelRequest.builder(request).messages(List.of(ChatMessage.user("short"))).build()) + .build(); + + interceptor.interceptModel(longRequest(), r -> { + dispatches.incrementAndGet(); + dispatched.set(r); + return response(); + }); + + assertThat(dispatches).hasValue(1); + assertThat(dispatched.get().getMessages()).hasSize(1); + assertThat(dispatched.get().getMessages().get(0).text()).isEqualTo("short"); + } + + @Test + void aCompactionWhichDoesNotHelpStillFails() { + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(20, 5)) + .compactor(request -> request) + .build(); + + assertThatThrownBy(() -> interceptor.interceptModel(request(), r -> { + dispatches.incrementAndGet(); + return response(); + })).isInstanceOf(ContextWindowBudgetExceededException.class); + + assertThat(dispatches).hasValue(0); + } + + @Test + void aCompactionWhichBreaksTheToolCallPairingIsRejected() { + ChatMessage call = ChatMessage.assistant("calling", List.of(new ToolCall("call-1", "search", "{}"))); + ChatMessage result = ChatMessage.tool("call-1", "search", "found it"); + ModelRequest request = ModelRequest.builder().messages(List.of(call, result)).build(); + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(5, 1)) + .compactor(original -> ModelRequest.builder(original).messages(List.of(call)).build()) + .build(); + + assertThatThrownBy(() -> interceptor.interceptModel(request, r -> { + dispatches.incrementAndGet(); + return response(); + })).isInstanceOf(ToolCallPairingValidator.PairingException.class) + .satisfies(failure -> assertThat(((ToolCallPairingValidator.PairingException) failure).violation().kind()) + .isEqualTo(ToolCallPairingValidator.Kind.MISSING_TOOL_RESULT)); + + assertThat(dispatches).hasValue(0); + } + + @Test + void theEstimatorIsReplaceable() { + AtomicInteger dispatches = new AtomicInteger(); + RequestTokenEstimator exact = (request, budget) -> new RequestTokenEstimate(1, 1, 1, budget.outputReserve(), + EstimationAccuracy.EXACT); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(20, 5)) + .estimator(exact) + .build(); + + interceptor.interceptModel(request(), r -> { + dispatches.incrementAndGet(); + return response(); + }); + + assertThat(dispatches).hasValue(1); + } + + @Test + void theBudgetBoundaryIsExact() { + RequestTokenEstimator fixed = (request, budget) -> new RequestTokenEstimate(50, 50, 0, 100, + EstimationAccuracy.EXACT); + AtomicInteger dispatches = new AtomicInteger(); + + ContextBudgetInterceptor fitting = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(200, 100)) + .estimator(fixed) + .build(); + fitting.interceptModel(request(), r -> { + dispatches.incrementAndGet(); + return response(); + }); + assertThat(dispatches).hasValue(1); + + ContextBudgetInterceptor tooSmall = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(199, 100)) + .estimator(fixed) + .build(); + assertThatThrownBy(() -> tooSmall.interceptModel(request(), r -> { + dispatches.incrementAndGet(); + return response(); + })).isInstanceOf(ContextWindowBudgetExceededException.class); + assertThat(dispatches).hasValue(1); + } + + private static ContextWindowBudget budget() { + return new ContextWindowBudget(1000, 100); + } + + private static ModelRequest request() { + return ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"), ChatMessage.assistant("hi"))) + .toolRegistry(new ToolRegistry(List.of(tool("search", "{\"type\":\"object\",\"properties\":{\"query\":{}}}")))) + .build(); + } + + private static ModelResponse response() { + return ModelResponse.of(ChatMessage.assistant("ok")); + } + + private static ModelRequest longRequest() { + List messages = new java.util.ArrayList<>(); + for (int index = 0; index < 20; index++) { + messages.add(ChatMessage.user("message number " + index + " with a fair amount of content")); + } + return ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(messages) + .toolRegistry(new ToolRegistry(List.of(tool("search", "{\"type\":\"object\",\"properties\":{\"query\":{}}}")))) + .build(); + } + + private static Tool tool(String name, String schema) { + return new Tool() { + + @Override + public ToolDefinition definition() { + return new ToolDefinition(name, "tool " + name, schema); + } + + @Override + public CompletionStage execute(ToolCall call, ToolContext context) { + return CompletableFuture.completedFuture(ToolResult.text("result")); + } + }; + } + +} diff --git a/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/extension/interceptor/LargeResultEvictionInterceptor.java b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/extension/interceptor/LargeResultEvictionInterceptor.java index 9bbc79610..a19c78e04 100644 --- a/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/extension/interceptor/LargeResultEvictionInterceptor.java +++ b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/extension/interceptor/LargeResultEvictionInterceptor.java @@ -151,15 +151,19 @@ private ToolCallResponse processLargeResult(ToolCallResponse response, String to String sanitizedId = sanitizeToolCallId(toolCallId); String filePath = LARGE_RESULTS_DIR + sanitizedId; - // Write content to filesystem via backend - if (backend != null) { - WriteResult writeResult = backend.write(filePath, content); - if (writeResult.getError() != null) { - // Log warning but continue with eviction message - log.warn("Failed to write large result to filesystem: {}", writeResult.getError()); - } - } else { + // Write content to filesystem via backend. A failed write must not produce a + // pointer to a file which was never written. + if (backend == null) { log.warn("Backend not configured, large result not persisted to filesystem"); + return evictionFailed(response, + "The tool result was too large and could not be persisted, because no filesystem backend is " + + "configured. The result is not available."); + } + WriteResult writeResult = backend.write(filePath, content); + if (writeResult.getError() != null) { + log.warn("Failed to write large result to filesystem: {}", writeResult.getError()); + return evictionFailed(response, "The tool result was too large and could not be persisted: " + + writeResult.getError() + ". The result is not available."); } // Extract first N lines as sample @@ -179,6 +183,26 @@ private ToolCallResponse processLargeResult(ToolCallResponse response, String to .build(); } + /** + * Builds the response of an eviction which did not persist the result. + * + *

+ * The response states that the result is not available. It never points to a file + * which was not written. + * @param response the original tool response + * @param reason the reason why the result could not be persisted + * @return the failed eviction response + */ + private ToolCallResponse evictionFailed(ToolCallResponse response, String reason) { + return ToolCallResponse.builder() + .content(reason) + .toolName(response.getToolName()) + .toolCallId(response.getToolCallId()) + .status(ToolCallResponse.ERROR_STATUS) + .metadata(response.getMetadata()) + .build(); + } + /** * Sanitize tool call ID for use in file paths. * Removes non-alphanumeric characters except underscore and hyphen. diff --git a/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/hook/messages/SummarizationRequestCompactor.java b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/hook/messages/SummarizationRequestCompactor.java new file mode 100644 index 000000000..e98df6b2f --- /dev/null +++ b/argi-orchestration/argi-agent-framework/src/main/java/io/github/agentic/ai/agent/orchestration/hook/messages/SummarizationRequestCompactor.java @@ -0,0 +1,77 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.orchestration.hook.messages; + +import io.github.agentic.ai.agent.budget.RequestCompactor; +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.orchestration.hook.summarization.SummarizationHook; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.workflow.graph.RunnableConfig; + +import java.util.List; +import java.util.Objects; + +/** + * Compresses a model request with the existing {@link SummarizationHook}. + * + *

+ * The adapter is the bridge between the neutral budget contract and the existing + * compression capability. It does not summarize anything itself: it hands the + * conversation to the hook and rebuilds the request with the messages the hook + * returned. + * + *

+ * The behaviour of the hook is not changed. When the hook decides that no summary + * is needed, or when the summary fails and the hook keeps the original history, the + * original request is returned unchanged and the budget gate decides what to do with + * it. + */ +public final class SummarizationRequestCompactor implements RequestCompactor { + + private final SummarizationHook hook; + + private final RunnableConfig config; + + /** + * Creates a compactor which summarizes without a run configuration. + * @param hook the summarization hook, must not be null + */ + public SummarizationRequestCompactor(SummarizationHook hook) { + this(hook, null); + } + + /** + * Creates a compactor which summarizes within a run. + * @param hook the summarization hook, must not be null + * @param config the run configuration, nullable + */ + public SummarizationRequestCompactor(SummarizationHook hook, RunnableConfig config) { + this.hook = Objects.requireNonNull(hook, "hook must not be null"); + this.config = config; + } + + @Override + public ModelRequest compact(ModelRequest request) { + Objects.requireNonNull(request, "request must not be null"); + AgentCommand command = this.hook.beforeModel(request.getMessages(), this.config); + List messages = command.getMessages(); + if (messages == null || messages.equals(request.getMessages())) { + return request; + } + return ModelRequest.builder(request).messages(messages).build(); + } + +} diff --git a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/hooks/messages/SummarizationRequestCompactorTest.java b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/hooks/messages/SummarizationRequestCompactorTest.java new file mode 100644 index 000000000..98be1dea5 --- /dev/null +++ b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/hooks/messages/SummarizationRequestCompactorTest.java @@ -0,0 +1,119 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.orchestration.hooks.messages; + +import io.github.agentic.ai.agent.interceptor.ModelRequest; +import io.github.agentic.ai.agent.orchestration.hook.messages.SummarizationRequestCompactor; +import io.github.agentic.ai.agent.orchestration.hook.summarization.SummarizationHook; +import io.github.agentic.ai.api.chat.ChatClient; +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ChatRequest; +import io.github.agentic.ai.api.chat.ChatResponse; +import io.github.agentic.ai.api.chat.FinishReason; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Verifies that the budget compactor reuses the existing summarization instead of + * implementing a second one. + */ +@DisplayName("SummarizationRequestCompactor") +class SummarizationRequestCompactorTest { + + @Test + @DisplayName("the conversation is summarized by the existing hook") + void summarizesTheConversationWithTheExistingHook() { + ChatClient client = mock(ChatClient.class); + when(client.call(any(ChatRequest.class))) + .thenReturn(CompletableFuture.completedFuture(new ChatResponse(null, null, + List.of(ChatMessage.assistant("Summary of the earlier conversation")), List.of(), FinishReason.STOP, + null, Map.of()))); + when(client.callBlocking(any(ChatRequest.class))).thenCallRealMethod(); + + SummarizationHook hook = SummarizationHook.builder() + .chatClient(client) + .maxTokensBeforeSummary(1) + .messagesToKeep(1) + .build(); + + ModelRequest request = ModelRequest.builder() + .systemMessage(ChatMessage.system("system prompt")) + .messages(conversation()) + .build(); + + ModelRequest compacted = new SummarizationRequestCompactor(hook).compact(request); + + assertThat(compacted.getMessages()).hasSizeLessThan(request.getMessages().size()); + assertThat(compacted.getMessages()) + .anySatisfy(message -> assertThat(message.text()).contains("Previous conversation summary")); + assertThat(compacted.getSystemMessage()).isEqualTo(request.getSystemMessage()); + } + + @Test + @DisplayName("a failed summary keeps the original request") + void keepsTheOriginalRequestWhenTheSummaryFails() { + ChatClient client = mock(ChatClient.class); + when(client.call(any(ChatRequest.class))) + .thenReturn(CompletableFuture.failedFuture(new IllegalStateException("the summary model is unavailable"))); + when(client.callBlocking(any(ChatRequest.class))).thenCallRealMethod(); + + SummarizationHook hook = SummarizationHook.builder() + .chatClient(client) + .maxTokensBeforeSummary(1) + .messagesToKeep(1) + .build(); + + ModelRequest request = ModelRequest.builder().messages(conversation()).build(); + + ModelRequest compacted = new SummarizationRequestCompactor(hook).compact(request); + + assertThat(compacted).isSameAs(request); + } + + @Test + @DisplayName("a conversation which does not need a summary is returned unchanged") + void keepsTheOriginalRequestWhenNoSummaryIsNeeded() { + ChatClient client = mock(ChatClient.class); + SummarizationHook hook = SummarizationHook.builder() + .chatClient(client) + .maxTokensBeforeSummary(100_000) + .messagesToKeep(1) + .build(); + + ModelRequest request = ModelRequest.builder().messages(conversation()).build(); + + ModelRequest compacted = new SummarizationRequestCompactor(hook).compact(request); + + assertThat(compacted).isSameAs(request); + } + + private static List conversation() { + return List.of(ChatMessage.user("first question with a fair amount of content"), + ChatMessage.assistant("first answer with a fair amount of content"), + ChatMessage.user("second question with a fair amount of content"), + ChatMessage.assistant("second answer with a fair amount of content"), + ChatMessage.user("third question with a fair amount of content")); + } + +} diff --git a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/LargeResultEvictionInterceptorTest.java b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/LargeResultEvictionInterceptorTest.java index cd215ee68..a49ddb602 100644 --- a/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/LargeResultEvictionInterceptorTest.java +++ b/argi-orchestration/argi-agent-framework/src/test/java/io/github/agentic/ai/agent/orchestration/interceptors/LargeResultEvictionInterceptorTest.java @@ -71,4 +71,59 @@ void evictionKeepsOriginalMetadata() { .containsEntry("errorCode", "SLOW_QUERY"); } + @Test + @DisplayName("a failed write does not produce a pointer to a file which was never written") + void failedWriteDoesNotClaimThatTheResultWasPersisted() { + FilesystemBackend backend = mock(FilesystemBackend.class); + when(backend.write(anyString(), anyString())).thenReturn(new WriteResult(null, "the disk is full", null)); + + LargeResultEvictionInterceptor interceptor = LargeResultEvictionInterceptor.builder() + .toolTokenLimitBeforeEvict(1) + .backend(backend) + .build(); + + ToolCallResponse original = ToolCallResponse.builder() + .content("x".repeat(500)) + .toolName("reportTool") + .toolCallId("call_1") + .status(ToolCallResponse.SUCCESS_STATUS) + .build(); + + ToolCallResponse failed = interceptor.interceptToolCall( + new ToolCallRequest("reportTool", "{}", "call_1", new HashMap<>()), r -> original); + + assertThat(failed.getStatus()).isEqualTo(ToolCallResponse.ERROR_STATUS); + assertThat(failed.getResult()).doesNotContain("saved in the filesystem").contains("could not be persisted"); + } + + @Test + @DisplayName("a failed write keeps the failure metadata of the original response") + void failedWriteKeepsTheOriginalMetadata() { + FilesystemBackend backend = mock(FilesystemBackend.class); + when(backend.write(anyString(), anyString())).thenReturn(new WriteResult(null, "the disk is full", null)); + + LargeResultEvictionInterceptor interceptor = LargeResultEvictionInterceptor.builder() + .toolTokenLimitBeforeEvict(1) + .backend(backend) + .build(); + + Map metadata = new HashMap<>(); + metadata.put(ToolCallResponse.FAILURE_KIND_METADATA_KEY, ToolCallResponse.FAILURE_KIND_EXECUTION); + ToolCallResponse original = ToolCallResponse.builder() + .content("x".repeat(500)) + .toolName("reportTool") + .toolCallId("call_1") + .status(ToolCallResponse.SUCCESS_STATUS) + .metadata(metadata) + .build(); + + ToolCallResponse failed = interceptor.interceptToolCall( + new ToolCallRequest("reportTool", "{}", "call_1", new HashMap<>()), r -> original); + + assertThat(failed.getStatus()).isEqualTo(ToolCallResponse.ERROR_STATUS); + assertThat(failed.getResult()).doesNotContain("saved in the filesystem").contains("could not be persisted"); + assertThat(failed.getMetadata()) + .containsEntry(ToolCallResponse.FAILURE_KIND_METADATA_KEY, ToolCallResponse.FAILURE_KIND_EXECUTION); + } + } From b5ad4018de81467a46d52c795b127189679d0cf6 Mon Sep 17 00:00:00 2001 From: konlue <2362586639@qq.com> Date: Tue, 3 Nov 2026 00:39:42 +0800 Subject: [PATCH 2/4] fix(agent): position-aware tool result pairing and documented budget ordering (issue #137 review) --- .../budget/ContextBudgetInterceptor.java | 12 +++ .../budget/ToolCallPairingValidator.java | 55 ++++++++---- .../budget/ContextBudgetInterceptorTest.java | 36 ++++++++ .../budget/ToolCallPairingValidatorTest.java | 85 +++++++++++++++++++ 4 files changed, 172 insertions(+), 16 deletions(-) create mode 100644 argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidatorTest.java diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java index b3bb6d975..ed0e87689 100644 --- a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptor.java @@ -69,6 +69,18 @@ public static Builder builder() { return new Builder(); } + /** + * Ordering contract: this interceptor must be registered **after** (i.e. more + * innermost than) every interceptor that rewrites the request. The budget is + * checked against the request as seen at this interceptor's position in the + * chain — an outer rewrite that grows the request after the budget check has + * passed would bypass the gate (a chain-level regression test documents this). + * Terminal, non-bypassable enforcement belongs to the runtime dispatch (see + * issue #137 discussion). + */ + public static final String ORDERING_CONTRACT = "register the context budget interceptor" + + " after every request-rewriting interceptor"; + @Override public ModelResponse interceptModel(ModelRequest request, ModelCallHandler handler) { Objects.requireNonNull(request, "request must not be null"); diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java index baab428e8..69cb4acf3 100644 --- a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidator.java @@ -19,13 +19,12 @@ import io.github.agentic.ai.api.chat.ChatRole; import io.github.agentic.ai.api.chat.ToolCall; +import java.util.ArrayList; import java.util.HashMap; -import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Set; /** * Validates that the tool calls of a conversation and their results stay paired. @@ -34,7 +33,9 @@ * Compression may remove a message. When it removes a tool result, the assistant * message which requested the tool call is left with an unanswered call; when it * removes the assistant message, the tool result is left without a call. Both break - * the conversation, so the budget gate refuses to dispatch such a request. + * the conversation, so the budget gate refuses to dispatch such a request. The + * validation is position-aware: a result must follow the assistant message which + * requested the call, and neither calls nor results may repeat. */ public final class ToolCallPairingValidator { @@ -59,7 +60,17 @@ public enum Kind { /** * The same tool call appears more than once. */ - DUPLICATE_TOOL_CALL + DUPLICATE_TOOL_CALL, + + /** + * The same tool result appears more than once. + */ + DUPLICATE_TOOL_RESULT, + + /** + * A tool result appears before the assistant message which requested the call. + */ + RESULT_BEFORE_CALL } /** @@ -90,14 +101,20 @@ public record Violation(Kind kind, String toolCallId) { public static Optional validate(List messages) { Objects.requireNonNull(messages, "messages must not be null"); Map callCounts = new HashMap<>(); - Set calls = new LinkedHashSet<>(); - for (ChatMessage message : messages) { + Map callFirstIndex = new HashMap<>(); + Map> resultIndexes = new HashMap<>(); + for (int index = 0; index < messages.size(); index++) { + ChatMessage message = messages.get(index); for (ToolCall call : message.toolCalls()) { if (call.id() == null) { continue; } callCounts.merge(call.id(), 1, Integer::sum); - calls.add(call.id()); + callFirstIndex.putIfAbsent(call.id(), index); + } + if (message.role() == ChatRole.TOOL && message.toolCallId() != null) { + resultIndexes.computeIfAbsent(message.toolCallId(), key -> new ArrayList<>()) + .add(index); } } for (Map.Entry call : callCounts.entrySet()) { @@ -105,19 +122,25 @@ public static Optional validate(List messages) { return Optional.of(new Violation(Kind.DUPLICATE_TOOL_CALL, call.getKey())); } } - Set results = new LinkedHashSet<>(); - for (ChatMessage message : messages) { - if (message.role() == ChatRole.TOOL && message.toolCallId() != null) { - results.add(message.toolCallId()); + for (Map.Entry> result : resultIndexes.entrySet()) { + if (result.getValue().size() > 1) { + return Optional.of(new Violation(Kind.DUPLICATE_TOOL_RESULT, result.getKey())); } } - for (String result : results) { - if (!calls.contains(result)) { - return Optional.of(new Violation(Kind.UNKNOWN_TOOL_RESULT, result)); + for (Map.Entry> result : resultIndexes.entrySet()) { + Integer callIndex = callFirstIndex.get(result.getKey()); + if (callIndex == null) { + return Optional.of(new Violation(Kind.UNKNOWN_TOOL_RESULT, result.getKey())); + } + // A result must follow the assistant message which requested the call: + // a compacted history that moves a result in front of its call would + // otherwise reach the provider. + if (result.getValue().get(0) < callIndex) { + return Optional.of(new Violation(Kind.RESULT_BEFORE_CALL, result.getKey())); } } - for (String call : calls) { - if (!results.contains(call)) { + for (String call : callCounts.keySet()) { + if (!resultIndexes.containsKey(call)) { return Optional.of(new Violation(Kind.MISSING_TOOL_RESULT, call)); } } diff --git a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java index 4d0cca4b5..5ae74d74c 100644 --- a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java +++ b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java @@ -15,6 +15,8 @@ */ package io.github.agentic.ai.agent.budget; +import io.github.agentic.ai.agent.interceptor.ModelInterceptor; +import io.github.agentic.ai.agent.interceptor.ModelCallHandler; import io.github.agentic.ai.agent.interceptor.ModelRequest; import io.github.agentic.ai.agent.interceptor.ModelResponse; import io.github.agentic.ai.api.chat.ChatMessage; @@ -56,6 +58,40 @@ void theEstimateCoversSystemMessagesToolSchemasAndTheOutputReserve() { assertThat(estimate.accuracy()).isEqualTo(EstimationAccuracy.APPROXIMATE); } + // issue #137 review: the budget must hold after interceptor rewrites, so the + // budget interceptor has to sit **after** every request-rewriting interceptor + // in the chain (innermost). This chain-level case registers [rewriter, budget]: + // a small request passes the rewriter, grows beyond the window, and must be + // refused before reaching the provider. + @Test + void budgetStillGuardsAfterARequestRewritingInterceptor() throws Exception { + AtomicInteger dispatches = new AtomicInteger(); + ModelInterceptor rewriter = new ModelInterceptor() { + public String getName() { + return "rewriter"; + } + + public ModelResponse interceptModel(ModelRequest request, ModelCallHandler next) { + // Rewrites the small incoming request into an oversized one. + return next.call(longRequest()); + } + }; + ContextBudgetInterceptor budget = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(20, 5)) + .build(); + // Registration order: the rewriter outermost, the budget innermost (the + // terminal position, right before the dispatch). + ModelCallHandler chain = request -> { + dispatches.incrementAndGet(); + return response(); + }; + ModelCallHandler budgetHandler = request -> budget.interceptModel(request, chain); + ModelCallHandler fullChain = request -> rewriter.interceptModel(request, budgetHandler); + + assertThatThrownBy(() -> fullChain.call(request())) + .isInstanceOf(ContextWindowBudgetExceededException.class); + assertThat(dispatches).hasValue(0); + } @Test void anOversizedRequestIsNeverDispatched() { AtomicInteger dispatches = new AtomicInteger(); diff --git a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidatorTest.java b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidatorTest.java new file mode 100644 index 000000000..a6e924b92 --- /dev/null +++ b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ToolCallPairingValidatorTest.java @@ -0,0 +1,85 @@ +/* + * Copyright 2026 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.github.agentic.ai.agent.budget; + +import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.ToolCall; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Position-aware pairing validation: a result must follow the assistant message + * which requested the call, and neither calls nor results may repeat (issue #137 + * review). + */ +class ToolCallPairingValidatorTest { + + @Test + void aWellPairedConversationPasses() { + List messages = List.of( + ChatMessage.assistant("use", List.of(new ToolCall("call-1", "lookup", "{}"))), + ChatMessage.tool("call-1", "lookup", "lookup result")); + + assertThat(ToolCallPairingValidator.validate(messages)).isEmpty(); + } + + @Test + void aResultBeforeItsCallIsRejected() { + List messages = List.of( + ChatMessage.tool("call-1", "lookup", "lookup result"), + ChatMessage.assistant("use", List.of(new ToolCall("call-1", "lookup", "{}")))); + + assertThat(ToolCallPairingValidator.validate(messages)) + .contains(new ToolCallPairingValidator.Violation( + ToolCallPairingValidator.Kind.RESULT_BEFORE_CALL, "call-1")); + } + + @Test + void aDuplicatedResultIsRejected() { + List messages = List.of( + ChatMessage.assistant("use", List.of(new ToolCall("call-1", "lookup", "{}"))), + ChatMessage.tool("call-1", "lookup", "lookup result"), + ChatMessage.tool("call-1", "lookup", "lookup result")); + + assertThat(ToolCallPairingValidator.validate(messages)) + .contains(new ToolCallPairingValidator.Violation( + ToolCallPairingValidator.Kind.DUPLICATE_TOOL_RESULT, "call-1")); + } + + @Test + void aMissingResultIsStillRejected() { + List messages = List.of( + ChatMessage.assistant("use", List.of(new ToolCall("call-1", "lookup", "{}")))); + + assertThat(ToolCallPairingValidator.validate(messages)) + .contains(new ToolCallPairingValidator.Violation( + ToolCallPairingValidator.Kind.MISSING_TOOL_RESULT, "call-1")); + } + + @Test + void anOrphanedResultIsStillRejected() { + List messages = List.of( + ChatMessage.tool("call-1", "lookup", "lookup result")); + + assertThat(ToolCallPairingValidator.validate(messages)) + .contains(new ToolCallPairingValidator.Violation( + ToolCallPairingValidator.Kind.UNKNOWN_TOOL_RESULT, "call-1")); + } + +} From 140693bff5491d6d7601c3cc882a8c628155bd4b Mon Sep 17 00:00:00 2001 From: konlue Date: Wed, 4 Nov 2026 23:43:03 +0800 Subject: [PATCH 3/4] fix(agent): account for an explicit maxTokens output demand in the budget estimate (issue #137 review) --- .../ApproximateRequestTokenEstimator.java | 20 +++++- .../budget/ContextBudgetInterceptorTest.java | 66 +++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java index 83c30cc58..8abf904c4 100644 --- a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java @@ -40,6 +40,12 @@ * the request carries an explicit tool selection, only the selected bindings are * counted; the dynamic bindings of the request are always counted because they are * disclosed next to the selected ones. + * + *

+ * When the request carries an explicit {@code maxTokens} output limit, the estimate + * accounts for that demand instead of the configured output reserve alone: the gate + * must refuse a request whose explicitly requested output does not fit into the + * window, even when the input itself is small. */ public final class ApproximateRequestTokenEstimator implements RequestTokenEstimator { @@ -73,7 +79,19 @@ public RequestTokenEstimate estimate(ModelRequest request, ContextWindowBudget b Objects.requireNonNull(request, "request must not be null"); Objects.requireNonNull(budget, "budget must not be null"); return new RequestTokenEstimate(systemTokens(request), messageTokens(request), toolSchemaTokens(request), - budget.outputReserve(), EstimationAccuracy.APPROXIMATE); + outputDemand(request, budget), EstimationAccuracy.APPROXIMATE); + } + + /** + * The output demand of the request: an explicit {@code maxTokens} option counts as + * the demanded output, floored by the configured output reserve. + */ + private int outputDemand(ModelRequest request, ContextWindowBudget budget) { + Integer maxTokens = request.getOptions() == null ? null : request.getOptions().maxTokens(); + if (maxTokens == null) { + return budget.outputReserve(); + } + return Math.max(budget.outputReserve(), maxTokens); } private int systemTokens(ModelRequest request) { diff --git a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java index 5ae74d74c..cd29ba034 100644 --- a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java +++ b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java @@ -20,6 +20,7 @@ import io.github.agentic.ai.agent.interceptor.ModelRequest; import io.github.agentic.ai.agent.interceptor.ModelResponse; import io.github.agentic.ai.api.chat.ChatMessage; +import io.github.agentic.ai.api.chat.GenerationOptions; import io.github.agentic.ai.api.chat.ToolCall; import io.github.agentic.ai.api.chat.ToolDefinition; import io.github.agentic.ai.api.tool.Tool; @@ -58,6 +59,71 @@ void theEstimateCoversSystemMessagesToolSchemasAndTheOutputReserve() { assertThat(estimate.accuracy()).isEqualTo(EstimationAccuracy.APPROXIMATE); } + // issue #137 review: an explicit maxTokens output limit is an output demand — + // a request whose requested completion alone exceeds the window must be refused + // even when its input is small. + @Test + void anExplicitOutputDemandBeyondTheWindowIsRefused() { + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(1000, 100)) + .build(); + ModelRequest demanded = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .options(GenerationOptions.builder().maxTokens(2000).build()) + .build(); + + assertThatThrownBy(() -> interceptor.interceptModel(demanded, r -> { + dispatches.incrementAndGet(); + return response(); + })).isInstanceOf(ContextWindowBudgetExceededException.class) + .satisfies(failure -> { + ContextWindowBudgetExceededException exceeded = (ContextWindowBudgetExceededException) failure; + assertThat(exceeded.dominantComponent()).isEqualTo(BudgetComponent.OUTPUT_RESERVE); + }); + + assertThat(dispatches).hasValue(0); + } + + // The same gate must not reject a request whose explicit output limit fits. + @Test + void anExplicitOutputDemandWithinTheWindowStillDispatches() { + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(1000, 100)) + .build(); + ModelRequest modest = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .options(GenerationOptions.builder().maxTokens(50).build()) + .build(); + + interceptor.interceptModel(modest, r -> { + dispatches.incrementAndGet(); + return response(); + }); + + assertThat(dispatches).hasValue(1); + } + + // The estimate itself accounts for the explicit output demand, floored by the + // configured reserve. + @Test + void theEstimateAccountsForAnExplicitMaxTokens() { + RequestTokenEstimate estimate = new ApproximateRequestTokenEstimator().estimate(request(), budget()); + + ModelRequest demanded = ModelRequest.builder() + .messages(List.of(ChatMessage.user("hello"))) + .options(GenerationOptions.builder().maxTokens(2000).build()) + .build(); + RequestTokenEstimate demandedEstimate = new ApproximateRequestTokenEstimator().estimate(demanded, budget()); + + assertThat(estimate.outputReserve()).isEqualTo(100); + assertThat(demandedEstimate.outputReserve()).isEqualTo(2000); + assertThat(demandedEstimate.totalTokens()).isEqualTo(demandedEstimate.inputTokens() + 2000); + } + // issue #137 review: the budget must hold after interceptor rewrites, so the // budget interceptor has to sit **after** every request-rewriting interceptor // in the chain (innermost). This chain-level case registers [rewriter, budget]: From 7e1033be73e55d2136fcb97fad300578d8d3e069 Mon Sep 17 00:00:00 2001 From: konlue Date: Thu, 5 Nov 2026 23:27:23 +0800 Subject: [PATCH 4/4] fix(agent): estimate the tool schemas the dispatch actually sends (issue #137 review) --- .../ApproximateRequestTokenEstimator.java | 32 +++++--- .../budget/ContextBudgetInterceptorTest.java | 73 +++++++++++++++++++ 2 files changed, 93 insertions(+), 12 deletions(-) diff --git a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java index 8abf904c4..be62d7c13 100644 --- a/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java +++ b/argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java @@ -36,10 +36,14 @@ * disclosed tool schemas and the reserved output. * *

- * The tool schemas are counted for the bindings the request actually discloses. When - * the request carries an explicit tool selection, only the selected bindings are - * counted; the dynamic bindings of the request are always counted because they are - * disclosed next to the selected ones. + * The tool schemas are counted for the bindings the request actually dispatches, + * matching the selection semantics of the dispatch path: the static and the dynamic + * registry are merged with the static binding winning a name conflict, and an + * explicit tool selection (including an empty one, which disables tools) filters the + * merged set. Schemas of bindings the dispatch will not send are not counted, so an + * empty selection or a large dynamic catalog cannot reject a request whose dispatch + * carries none of them. Tool descriptions are not counted: they feed the auxiliary + * selection call, never the dispatched request. * *

* When the request carries an explicit {@code maxTokens} output limit, the estimate @@ -124,18 +128,22 @@ else if (part.data() != null) { } private int toolSchemaTokens(ModelRequest request) { - Set selection = request.hasToolSelection() ? new HashSet<>(request.getTools()) : null; - int total = 0; + // Merge both registries with the static binding winning a name conflict, then + // apply the explicit selection — the same semantics as the dispatch path. + Map disclosed = new java.util.LinkedHashMap<>(); for (ToolDefinition definition : request.getToolRegistry().definitions()) { - if (selection == null || selection.contains(definition.name())) { - total += definitionTokens(definition); - } + disclosed.putIfAbsent(definition.name(), definition); } for (ToolDefinition definition : request.getDynamicToolRegistry().definitions()) { - total += definitionTokens(definition); + disclosed.putIfAbsent(definition.name(), definition); } - for (Map.Entry description : request.getToolDescriptions().entrySet()) { - total += tokens(description.getKey()) + tokens(description.getValue()); + if (request.hasToolSelection()) { + Set selection = new HashSet<>(request.getTools()); + disclosed.values().removeIf(definition -> !selection.contains(definition.name())); + } + int total = 0; + for (ToolDefinition definition : disclosed.values()) { + total += definitionTokens(definition); } return total; } diff --git a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java index cd29ba034..063f0f931 100644 --- a/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java +++ b/argi-agent/argi-agent-api/src/test/java/io/github/agentic/ai/agent/budget/ContextBudgetInterceptorTest.java @@ -124,6 +124,79 @@ void theEstimateAccountsForAnExplicitMaxTokens() { assertThat(demandedEstimate.totalTokens()).isEqualTo(demandedEstimate.inputTokens() + 2000); } + // issue #137 review: the tool schemas must match what the dispatch actually sends. + // The empty explicit selection disables tools, so a large dynamic catalog must not + // count against the request. + @Test + void anEmptyToolSelectionDoesNotCountDynamicSchemas() { + AtomicInteger dispatches = new AtomicInteger(); + ContextBudgetInterceptor interceptor = ContextBudgetInterceptor.builder() + .budget(new ContextWindowBudget(1000, 100)) + .build(); + ModelRequest disabled = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .toolRegistry(new ToolRegistry(List.of())) + .dynamicToolRegistry(new ToolRegistry(List.of(tool("catalog", "{\"type\":\"object\",\"properties\":{" + + "field1\":{\"type\":\"string\"},field2:{\"type\":\"string\"}}}")))) + .tools(List.of()) + .build(); + RequestTokenEstimate estimate = new ApproximateRequestTokenEstimator().estimate(disabled, budget()); + assertThat(estimate.toolSchemaTokens()).isZero(); + + interceptor.interceptModel(disabled, r -> { + dispatches.incrementAndGet(); + return response(); + }); + assertThat(dispatches).hasValue(1); + } + + // An explicit selection filters the dynamic registry exactly like the dispatch does. + @Test + void anExplicitSelectionFiltersDynamicSchemas() { + ModelRequest selected = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .toolRegistry(new ToolRegistry(List.of(tool("lookup", "{\"type\":\"object\"}")))) + .dynamicToolRegistry(new ToolRegistry(List.of(tool("catalog", "{\"type\":\"object\",\"properties\":{" + + "field1\":{\"type\":\"string\"},field2:{\"type\":\"string\"}}}")))) + .tools(List.of("lookup")) + .build(); + RequestTokenEstimate estimate = new ApproximateRequestTokenEstimator().estimate(selected, budget()); + + ModelRequest staticOnly = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .toolRegistry(new ToolRegistry(List.of(tool("lookup", "{\"type\":\"object\"}")))) + .build(); + RequestTokenEstimate staticEstimate = new ApproximateRequestTokenEstimator().estimate(staticOnly, budget()); + + assertThat(estimate.toolSchemaTokens()).isEqualTo(staticEstimate.toolSchemaTokens()); + } + + // A dynamic binding shadowed by a static name is dispatched once — the schema is + // counted once, not twice. + @Test + void aDynamicSchemaShadowedByAStaticNameIsCountedOnce() { + ModelRequest shadowed = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .toolRegistry(new ToolRegistry(List.of(tool("lookup", "{\"type\":\"object\"}")))) + .dynamicToolRegistry(new ToolRegistry(List.of(tool("lookup", "{\"type\":\"object\",\"properties\":{" + + "field1\":{\"type\":\"string\"},field2:{\"type\":\"string\"}}}")))) + .build(); + RequestTokenEstimate estimate = new ApproximateRequestTokenEstimator().estimate(shadowed, budget()); + + ModelRequest staticOnly = ModelRequest.builder() + .systemMessage(ChatMessage.system("system instructions for the agent")) + .messages(List.of(ChatMessage.user("hello"))) + .toolRegistry(new ToolRegistry(List.of(tool("lookup", "{\"type\":\"object\"}")))) + .build(); + RequestTokenEstimate staticEstimate = new ApproximateRequestTokenEstimator().estimate(staticOnly, budget()); + + assertThat(estimate.toolSchemaTokens()).isEqualTo(staticEstimate.toolSchemaTokens()); + } + // issue #137 review: the budget must hold after interceptor rewrites, so the // budget interceptor has to sit **after** every request-rewriting interceptor // in the chain (innermost). This chain-level case registers [rewriter, budget]: