Skip to content

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

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

konlue wants to merge 4 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
Member

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.

@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.

@yuluo-yx

yuluo-yx commented Nov 4, 2026

Copy link
Copy Markdown
Member

[P2] The budget gate can dispatch a request whose requested output alone exceeds the context window.

ApproximateRequestTokenEstimator.java:75 always uses budget.outputReserve() and does not reconcile it with request.getOptions().maxTokens(). Reproduced with ContextWindowBudget(1000, 100), a short user message and maxTokens=2000: the estimate is 102 tokens and the handler is called, even though the requested completion alone exceeds the entire window.

Please validate the actual output limit against the reserved budget, account for it in the estimate, or clamp it before dispatch. Otherwise, even with the interceptor registered last, an oversized request still reaches the provider.

Validation: targeted offline Java probe against this PR head; the full reactor test suite was not run.

@konlue

konlue commented Nov 4, 2026

Copy link
Copy Markdown
Contributor Author

Fixed in 140693b (+85/−1). The estimator now treats an explicit maxTokens option as an output demand: ApproximateRequestTokenEstimator computes max(budget.outputReserve(), maxTokens) when the request carries the option, so the gate measures the requested completion rather than assuming the configured reserve. Reproducing your probe — ContextWindowBudget(1000, 100), a short message and maxTokens=2000 — the estimate becomes 2002 tokens, the dispatch is refused with a ContextWindowBudgetExceededException whose dominant component is OUTPUT_RESERVE, and the handler is never called. A request whose explicit limit fits (e.g. maxTokens=50) still dispatches, and requests without the option keep the previous reserve-based estimate.

Three new cases in ContextBudgetInterceptorTest cover the refusal, the pass-through, and the estimate decomposition (11/11 green in the module; checkstyle clean).

@yuluo-yx

yuluo-yx commented Nov 5, 2026

Copy link
Copy Markdown
Member

[P2] Dynamic tool schemas ignore explicit tool selection (line 134). Reproduced: .tools(List.of()) plus an 8,000-character dynamic schema rejects a short request against a 1,000-token window, although both engines dispatch zero tools. Please match dispatch selection and static-over-dynamic deduplication when counting schemas.

@konlue

konlue commented Nov 5, 2026

Copy link
Copy Markdown
Contributor Author

Fixed in 7e1033b (+93/−12). ApproximateRequestTokenEstimator.toolSchemaTokens now mirrors the dispatch semantics (AgentLlmNode.buildChatRequest / ReActAgent.effectiveToolSet): the static and dynamic registries are merged with the static binding winning a name conflict (putIfAbsent), and an explicit tool selection — including an empty one — filters the merged set before counting. Reproducing your probe — an empty selection plus a large dynamic schema — the dynamic schemas no longer count against the request, so it dispatches; a selection that keeps only static tools counts exactly those; and a dynamic binding shadowed by a static name is counted once instead of twice.

One deliberate scope adjustment in the same area, stated explicitly: tool descriptions are no longer counted at all. They never enter the dispatched ChatRequest — they feed the auxiliary selection call of ToolSelectionInterceptor, a different model call with its own context — so counting them here overestimated every request and was the same over-rejection class as the dynamic-catalog case. The javadoc documents the new coverage ("the bindings the request actually dispatches") and this exclusion.

Three new cases in ContextBudgetInterceptorTest (empty selection ignores dynamic schemas and dispatches; explicit selection filters them; shadowed dynamic schema counted once) — 14/14 green in the module, checkstyle clean.

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