feat(agent): gate the final model request on a context window budget (issue #137)

argi-agent/argi-agent-api/src/main/java/io/github/agentic/ai/agent/budget/ApproximateRequestTokenEstimator.java

@@ -0,0 +1,169 @@
+/*
+ * 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.
+ *
+ * <p>
+ * 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.
+ *
+ * <p>
+ * 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.
+ *
+ * <p>
+ * 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 {
+
+	/**
+	 * 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),
+				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) {
+		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) {
+		// Merge both registries with the static binding winning a name conflict, then
+		// apply the explicit selection — the same semantics as the dispatch path.
+		Map<String, ToolDefinition> disclosed = new java.util.LinkedHashMap<>();
+		for (ToolDefinition definition : request.getToolRegistry().definitions()) {
+			disclosed.putIfAbsent(definition.name(), definition);
+		}
+		for (ToolDefinition definition : request.getDynamicToolRegistry().definitions()) {
+			disclosed.putIfAbsent(definition.name(), definition);
+		}
+		if (request.hasToolSelection()) {
+			Set<String> 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;
+	}
+
+	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;
+	}
+
+}

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
+
+}

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.
+ *
+ * <p>
+ * 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);
+	}
+
+}