Skip to content

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

Open
konlue wants to merge 2 commits into
agentic-ai-java:mainfrom
konlue:feat/issue137-context-window-budget
Open

konlue wants to merge 2 commits into
agentic-ai-java:mainfrom
konlue:feat/issue137-context-window-budget

Conversation

@konlue

@konlue konlue commented Nov 2, 2026

Copy link
Copy Markdown
Contributor

Describe what this PR does / why we need it

Fixes #137: the neutral agent API had no notion of a context window budget. TokenCounter.countTokens(List<ChatMessage>) only measures the conversation, while the request which is finally dispatched carries more than the conversation - ReActAgent.invokeModel builds new ChatRequest(composeMessages(request), effective.definitions(), request.getOptions(), request.getAttributes()), so the system message, the tool schemas and the reserved output belong to the request but stay outside the measured scope. SummarizationHook and ContextEditingInterceptor each trigger on their own message threshold and both continue when they cannot compress anything, so nothing keeps an oversized request from reaching the provider, and LargeResultEvictionInterceptor claimed that a result was written to the filesystem even when the write failed.

Does this pull request fix one issue?

Fixes #137

Describe how you did it

  • Added the neutral budget contract io.github.agentic.ai.agent.budget in argi-agent-api: ContextWindowBudget for the context window and the reserved output, RequestTokenEstimate for the decomposition into system, conversation, tool schema and reserved output with an EstimationAccuracy that tells exact and conservative counts apart, RequestTokenEstimator with the default ApproximateRequestTokenEstimator, BudgetComponent and BudgetViolation for a stable classification, and ContextWindowBudgetExceededException for an explicit failure;
  • Added ContextBudgetInterceptor, a neutral ModelInterceptor which is opt-in through the existing interceptor chain: it estimates the request, asks the configured RequestCompactor to shrink an oversized request, measures the compacted request again, validates the tool call pairing and only then dispatches; a request which still does not fit fails with ContextWindowBudgetExceededException and never reaches the provider;
  • Reused the existing compression instead of rebuilding it: the framework adapter SummarizationRequestCompactor delegates to the existing SummarizationHook and rebuilds the request from the messages the hook returned; the hook keeps its own behaviour, so a failed summary returns the original history and the gate then fails explicitly instead of continuing with a broken context;
  • Fixed the write failure of the large result eviction: when the filesystem write fails, the interceptor no longer returns a pointer to a file which was never written, but an explicit failure which keeps the metadata of the original response;
  • Kept the existing behaviour intact: no existing signature, threshold or default changed, the new path is opt-in, and the neutral package depends on argi-api and argi-runtime-api only.

How was it tested / verified?

New ContextBudgetInterceptorTest in argi-agent-api (offline, no model and no provider):

  1. the estimate covers the system message, the tool schemas and the reserved output;
  2. an oversized request fails and the dispatch counter stays at zero;
  3. a compacted request is measured again and then dispatched;
  4. a compaction which does not help still fails without a dispatch;
  5. a compaction which breaks the tool call pairing is rejected without a dispatch;
  6. the estimator is replaceable, an exact estimator changes the decision;
  7. the boundary is exact, a request which exactly fills the budget is dispatched and one token more is rejected.

New SummarizationRequestCompactorTest and two additional cases in LargeResultEvictionInterceptorTest in argi-agent-framework: the compactor summarizes with the existing hook and returns the original request when the summary fails or is not needed, and a failed filesystem write produces a failure which contains no file pointer and no "saved in the filesystem" wording while keeping the metadata of the original response.

Green after the change: argi-agent-api 41/41, argi-agent-framework 1341/1341, argi-agent-core 138/138, argi-workflow-graph-core 551/551; spring-javaformat:validate and checkstyle:check report no violation for the two touched modules, codespell and the newline check pass.

Additional notes

  • Scope is the budget contract, the gate and the correctness of the write failure; the chain executor, the declarative plugin model and the engine wiring stay with their own issues, and the summarization and the context editing are not migrated into plugins here;
  • The change is additive: eleven new types in argi-agent-api, one new adapter and one behavioural fix in argi-agent-framework, and no change to an existing public signature, event, configuration or persistence format;
  • The builder of LargeResultEvictionInterceptor always supplies a default filesystem backend when eviction is enabled, so the branch for a missing backend is defensive only and the test covers the failed write path.

@yuluo-yx yuluo-yx left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The dispatch boundary and tool-history validation need tightening.

Objects.requireNonNull(handler, "handler must not be null");
RequestTokenEstimate estimate = estimator.estimate(request, budget);
if (estimate.totalTokens() <= budget.inputTokenLimit()) {
return handler.call(request);

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This checks the request before calling the next interceptor, not after all rewrites. Registering [budget, requestRewriter] lets a small request pass, then grow beyond the window and reach the provider unchecked. Please enforce the budget at the terminal dispatch boundary, or enforce and document the required ordering, with a chain-level regression case.

return Optional.of(new Violation(Kind.DUPLICATE_TOOL_CALL, call.getKey()));
}
}
Set<String> results = new LinkedHashSet<>();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sets only check whether IDs exist. Both [toolResult, assistantCall] and [assistantCall, toolResult, toolResult] pass, so invalid compacted history can still reach the provider. Please validate result ordering and reject duplicate results as well as missing ones.

@konlue

konlue commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the review — both points are addressed (b5ad401, +96/−17):

  1. Pairing validation is now position-aware (ToolCallPairingValidator): the scan records where each call and each result first appear, and two new violation kinds are reported — RESULT_BEFORE_CALL (a result whose first appearance precedes the assistant message requesting the call, e.g. [toolResult, assistantCall]) and DUPLICATE_TOOL_RESULT (the same result id twice, e.g. [assistantCall, toolResult, toolResult]). The three existing checks (missing result, orphaned result, duplicate call) are unchanged and covered by a new component-level ToolCallPairingValidatorTest (5 cases).

  2. Budget ordering: the interceptor's javadoc now states the ordering contract explicitly — the budget interceptor must be registered after (more innermost than) every request-rewriting interceptor, because the check applies to the request as seen at its chain position. A chain-level regression case was added (budgetStillGuardsAfterARequestRewritingInterceptor): a small request passes an outer rewriter, is rewritten beyond the window, and is refused before reaching the provider (dispatches = 0). 8/8 green in ContextBudgetInterceptorTest.

On terminal enforcement: making the budget non-bypassable at the runtime dispatch belongs to the lifecycle execution chain (issue #299's "non-bypassable runtime checks") — happy to take that as a follow-up once the contract PR lands.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[P1][T34] 统一完整模型请求的上下文窗口预算与压缩策略

2 participants