fix(agents): share the invocation cost manager so maxLlmCalls spans the whole run#484
Merged
kalenkevich merged 2 commits intoJul 15, 2026
Conversation
…he whole run `maxLlmCalls` is meant to cap the total number of LLM calls for an entire invocation, which "can contain one or multiple agent calls" per the InvocationContext docs. However, the LLM-call counter lives on a per-context `InvocationCostManager` that the field initializer always constructed fresh, and the constructor never adopted the parent's. Because a new InvocationContext is created for every sub-agent, agent transfer and loop iteration (createInvocationContext / createBranchCtxForSubAgent), the counter reset to 0 at each agent boundary — so the cap was effectively enforced per agent invocation, not per run. Concretely, a LoopAgent wrapping a non-terminating LlmAgent never hit the default limit of 500: each iteration built a new context with a fresh counter, so the "Max number of llm calls limit ... exceeded" guard never fired, allowing an unbounded number of LLM calls (runaway cost). Fix: adopt the parent invocation's cost manager in the constructor when present. Child contexts already copy every other parent field via the context spread / instance pass, so reusing the cost manager keeps a single shared counter for the whole invocation; a brand-new invocation (Runner, A2A) still starts a fresh one. This matches adk-python, where the cost manager is shared across agent calls. Adds regression tests proving the counter accumulates across child and grandchild contexts while remaining isolated between separate invocations.
kalenkevich
approved these changes
Jul 14, 2026
| @@ -0,0 +1,100 @@ | |||
| /** | |||
| * @license | |||
| * Copyright 2025 Google LLC | |||
| } as unknown as Session; | ||
| } | ||
|
|
||
| describe('InvocationContext LLM-call cost tracking', () => { |
Collaborator
There was a problem hiding this comment.
could you please add a test for your change as well? Thanks!
Contributor
Author
There was a problem hiding this comment.
Done! Added an end-to-end regression test (enforces maxLlmCalls across a real multi-iteration LoopAgent run) that drives a LoopAgent plus its sub-agent through the real createInvocationContext path. It fails on the pre-fix code (the run goes unbounded) and passes with the shared cost manager. Also bumped the copyright year to 2026. Thanks for the review!
Collaborator
There was a problem hiding this comment.
Sorry, but I does not see changes where you create an instance of InvocationContext and passing the InvocationCostManager. Can you add a tests for that as well? Thanks
Collaborator
There was a problem hiding this comment.
my bad, there is a test for that. Sorry!
Add an end-to-end regression test that drives a LoopAgent + sub-agent through the real createInvocationContext path, proving maxLlmCalls now bounds the whole run: it fails on the pre-fix code (the loop runs unbounded) and passes with the shared cost manager. The existing tests construct child contexts by hand; this one exercises the actual agent-run code path the fix targets. Also bump the test file's copyright year to 2026 to match the repo convention for files added in 2026.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
maxLlmCalls(default500) is documented as a limit on the total number of LLM calls for a given run, andInvocationContextis explicitly described as an invocation that "can contain one or multiple agent calls." In practice, though, the cap resets at every agent boundary, so it does not bound a multi-agent run at all.The LLM-call counter lives on a per-context
InvocationCostManager. The field was initialized eagerly (= new InvocationCostManager()) and the constructor never adopted the parent's, so every newInvocationContextstarts the counter back at 0. Because a fresh context is created for each sub-agent, agent transfer, and loop iteration (BaseAgent.createInvocationContext,createBranchCtxForSubAgent), the limit is effectively enforced per agent invocation rather than per run.Impact
A
LoopAgent(defaultmaxIterationsis unbounded) wrapping a non-terminatingLlmAgentnever trips the defaultmaxLlmCalls = 500: each iteration builds a new context whose counter resets to ~1, theMax number of llm calls limit ... exceededguard never fires, and the run makes an unbounded number of LLM calls — exactly the runaway-cost scenariomaxLlmCallsexists to prevent. The same reset happens on agent transfer.Fix
Adopt the parent invocation's cost manager in the
InvocationContextconstructor when one is present. Child contexts already copy every other parent field (via the context spread increateInvocationContextand the instance pass increateBranchCtxForSubAgent), so reusing the cost manager keeps a single shared counter for the whole invocation. A brand-new invocation (from theRunneror A2A executor) still gets a fresh counter.This matches adk-python, where the invocation cost manager is shared across agent calls (via
model_copy, which keeps the same manager reference).The change is intentionally minimal and touches no public API:
InvocationCostManagerstays private, and no new fields are added to the exportedInvocationContextParams.Test plan
core/test/agents/invocation_context_test.ts:npx vitest run --project unit:core): 2044 passing.tsc --noEmit, build (declaration emit), ESLint, and Prettier all clean.