perf(investigation): cut context-budget hot-path token recomputation - #4165
Conversation
Greptile code reviewThis repo uses Greptile for automated review. Before merge, aim for Confidence Score: 5/5 with zero unresolved review threads — see CONTRIBUTING.md. Run a review — add a PR comment with: Give it ~5-10 minutes (sometimes longer) for results, then fix feedback and re-trigger until you reach Confidence Score: 5/5. Optional: automate with the greploop skill. |
Greptile SummaryThis PR reduces repeated work in the investigation context-budget path. The main changes are:
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (3): Last reviewed commit: "feat(context_budget): introduce system_a..." | Re-trigger Greptile |
Davidson3556
left a comment
Davidson3556
left a comment
There was a problem hiding this comment.
Solid win on the hoist. Serializing 100+ tool schemas inside the trim loop was genuinely quadratic, and the subtraction in _truncate_largest_message checks out as equivalent to the old list comprehension. No behavior change that I can find.
Three small things before merge, all in the same direction (less new surface, not more):
system_tools_overhead_tokenshas no production caller.enforce_context_budgetuses the private helpers directly andagent.py:148doesn't pass it, so its only user is a test. It's also a footgun: pass it alongsidetools=andtoolsgets silently ignored.tool_sourceincore/execution.pyis orphaned now that both call sites moved, and_runtime_tool_sourcereimplements it. AGENTS.md asks for one canonical path after a refactor._refresh_message_token_estimatescan only ever take the recompute branch, sincetrim_lowest_value_tool_pairalways deletes at least one message. Worth collapsing, and worth editing the PR description: the per-message cache doesn't help the trim path, the hoisted overhead is the whole win.
Also, the tests hardcode 0.50 in four places instead of importing _TOKENS_PER_CHAR, and most of them recompute the production formula to build the expected value. The dump-count test is the useful one, keep that.
| def estimate_message_tokens( | ||
| messages: list[dict[str, Any]], | ||
| *, | ||
| system: str | None = None, | ||
| tools: list[dict[str, Any]] | None = None, | ||
| system_tools_overhead_tokens: int | None = None, | ||
| ) -> int: |
There was a problem hiding this comment.
Test-only parameter, and passing it alongside tools= silently ignores tools. Drop it:
| def estimate_message_tokens( | |
| messages: list[dict[str, Any]], | |
| *, | |
| system: str | None = None, | |
| tools: list[dict[str, Any]] | None = None, | |
| system_tools_overhead_tokens: int | None = None, | |
| ) -> int: | |
| def estimate_message_tokens( | |
| messages: list[dict[str, Any]], | |
| *, | |
| system: str | None = None, | |
| tools: list[dict[str, Any]] | None = None, | |
| ) -> int: |
| if system_tools_overhead_tokens is None: | ||
| system_tools_overhead_tokens = _system_and_tools_tokens(system, tools) | ||
| return _estimate_messages_tokens(messages) + system_tools_overhead_tokens |
There was a problem hiding this comment.
And the body collapses to one line once the override is gone:
| if system_tools_overhead_tokens is None: | |
| system_tools_overhead_tokens = _system_and_tools_tokens(system, tools) | |
| return _estimate_messages_tokens(messages) + system_tools_overhead_tokens | |
| return _estimate_messages_tokens(messages) + _system_and_tools_tokens(system, tools) |
| message_tokens, total_message_tokens = _refresh_message_token_estimates( | ||
| messages, | ||
| current_estimates=message_tokens, | ||
| ) |
There was a problem hiding this comment.
trim_lowest_value_tool_pair ends in del messages[start:end] with end > start guaranteed, so length always shrinks and the equal-length branch of _refresh_message_token_estimates is dead. Call the estimator directly and delete the helper:
| message_tokens, total_message_tokens = _refresh_message_token_estimates( | |
| messages, | |
| current_estimates=message_tokens, | |
| ) | |
| message_tokens, total_message_tokens = _message_token_estimates(messages) |
| msg[key] = True | ||
|
|
||
|
|
||
| def _runtime_tool_source(tool_by_name: dict[str, Any], tool_name: str) -> str: |
There was a problem hiding this comment.
This duplicates tool_source in core/execution.py, which now has no callers left outside the lazy re-export in core/__init__.py. Rather than reimplement the getattr(..., "source", "unknown") fallback here, prefer changing the shared one to take the map you already build:
def tool_source(tools: Mapping[str, RuntimeTool], tool_name: str) -> str:
tool = tools.get(tool_name)
return str(getattr(tool, "source", "unknown")) if tool else "unknown"then call tool_source(tool_by_name, tc.name) and drop _runtime_tool_source. Keeps one canonical path per AGENTS.md.
|
@Devesh36 I will take a look at the PR during a day today as well. I am interested to review |
yeah i'll take a look at davidson review |
…undant function and improving overhead calculation
|
@greptile-apps review again |
There was a problem hiding this comment.
Pull request overview
This PR optimizes the investigation context-budget enforcement hot path by separating message token estimation from system/tools overhead, precomputing invariant overhead once per enforcement call, and updating only the impacted message token estimate during truncation. It also improves evidence source lookup during gather-evidence by replacing repeated linear scans with an O(1) name→tool map.
Changes:
- Refactored
core/context_budget.pytoken estimation to precompute system+tools overhead once and maintain per-message token estimates during truncation. - Updated tool source lookup to use a prebuilt
tool_by_namemapping (and adjustedtool_sourceto accept mappings). - Added tests in
tests/core/test_context_budget.pyto characterize overhead inclusion and verify tool-schema serialization behavior.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| core/context_budget.py | Splits message vs. fixed overhead token estimation and avoids repeated tool-schema overhead recomputation in the budget loop. |
| core/execution.py | Changes tool_source to accept a name→tool mapping for O(1) lookup. |
| tools/investigation/stages/gather_evidence/agent.py | Builds a tool_by_name dict once and uses it for evidence entry source attribution. |
| tests/core/test_context_budget.py | Adds tests covering overhead inclusion and serialization frequency assumptions. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def counting_dumps(value: object, *args: object, **kwargs: object) -> str: | ||
| nonlocal schema_dump_calls | ||
| if isinstance(value, dict) and value.get("type") == "function": | ||
| schema_dump_calls += 1 | ||
| return original_dumps(value, *args, **kwargs) | ||
|
|
||
| with patch("core.context_budget.json.dumps", side_effect=counting_dumps): | ||
| enforce_context_budget(messages, tools=tools, ceiling=ceiling) | ||
|
|
||
| assert schema_dump_calls == len(tools) |
| def _message_token_estimates(messages: list[dict[str, Any]]) -> tuple[list[int], int]: | ||
| tokens = [_message_token_estimate(message) for message in messages] | ||
| return tokens, sum(tokens) |
| elif isinstance(block, str): | ||
| total += int(len(block) * _TOKENS_PER_CHAR) | ||
| return total | ||
|
|
There was a problem hiding this comment.
I can see json.dumps(block, default=str)) * _TOKENS_PER_CHAR and I feel it should be improved
I would like to have this method like:
def _message_token_estimate(message: dict[str, Any]) -> int:
content = message.get("content", "")
if isinstance(content, str):
chars = len(content)
elif isinstance(content, list):
chars = _sum_text_chars(content)
else:
chars = 0
return int(chars * _TOKENS_PER_CHAR)
There was a problem hiding this comment.
Thanks for the suggestion — I looked at this closely, but I don’t think _sum_text_chars is a safe drop-in here.
_sum_text_chars only counts strings under content/text keys (it’s built for the truncation path). A tool_use block like {"type": "tool_use", "id": "t1", "name": "noop", "input": {}} has neither key, so _sum_text_chars returns 0, while the serialized block is ~60 chars that the provider actually bills for.
Even for text blocks, the provider sees the full JSON envelope, so json.dumps is the more faithful upper bound for token estimation.
There was a problem hiding this comment.
Could you please think how to avoid calling json.dumps each time in for block in content:
| for schema in tools: | ||
| total += int(len(json.dumps(schema, default=str)) * _TOKENS_PER_CHAR) | ||
| return total | ||
|
|
There was a problem hiding this comment.
_system_and_tools_tokens calls json.dumps for every tool schema on each enforce_context_budget call. We need to move it out of the loop
Prefer computing that overhead once at schema-build time and passing the int in:
tool_schemas = llm.tool_schemas(tools)
tools_tokens = sum(
int(len(json.dumps(schema, default=str)) * _TOKENS_PER_CHAR)
for schema in tool_schemas
)
enforce_context_budget(
messages,
system=system,
tools_tokens=tools_tokens, # precomputed once per investigation
ceiling=ceiling,
)
Then the budget helper only needs system length:
def _system_tokens(system: str | None) -> int:
return int(len(system) * _TOKENS_PER_CHAR) if system else 0
But please write tests before refactoring, do these changes, and after that that all tests are passed
There was a problem hiding this comment.
Agreed — this was the right call.
I added a public system_and_tools_overhead(system, tools) helper and a fixed_overhead_tokens param on enforce_context_budget. The investigation agent now precomputes full_overhead and system_only_overhead once before the ReAct loop and passes the right value each iteration (including the force-conclusion path where tools are dropped).
I also applied the same pattern in core/agent/react_loop.py, which had the same per-iteration re-serialization issue.
Added a test that asserts zero tool-schema json.dumps calls when fixed_overhead_tokens is passed. All existing tests still pass.
|
mmm @YauhenBichel anything else ? |
|
@Devesh36 please use TDD: before changes check that tests cover all functionality, which you are refactoring, after that do refactoring, and be sure that all tests are green after that |
…r optimized token management
| if fixed_overhead_tokens is None: | ||
| fixed_overhead_tokens = system_and_tools_overhead(system, tools) | ||
| message_tokens, total_message_tokens = _message_token_estimates(messages) | ||
| while (total_message_tokens + fixed_overhead_tokens) > ceiling: |
yes sirrr |
|
Followed TDD for this refactor — existing characterization tests pinned the current behavior before the change (trimming, truncation, system/tools overhead), kept them green through the refactor, and added a new test for the precomputed overhead path. Ran uv run pytest tests/core/test_context_budget.py tests/agent/test_investigation.py — 43 passed. |
|
@greptile review |
|
@Devesh36 please be free to merge, also please check after merging manually that all is working. Good work, man |
|
🐉 Legend says enough merged PRs and you ascend. @Devesh36 is dangerously close. 🌤️
👋 Join us on Discord - OpenSRE : hang out, contribute, or hunt for features and issues. Everyone's welcome. |
This PR optimizes the investigation context-budget hot path by removing repeated token estimation work inside the ReAct loop, while preserving existing behavior.
What changed
core/context_budget.py:system + toolsoverhead once per enforcement passtools/investigation/stages/gather_evidence/agent.pyby using a prebuilttool_by_namemap instead of repeated linear scans.tests/core/test_context_budget.pyto characterize behavior and validate overhead/token-estimation paths.Why we did this
Investigation runs a ReAct loop where
enforce_context_budgetis called repeatedly. Before this change, token estimation frequently re-walked message history and re-serialized large tool schema payloads (viajson.dumps) across loop iterations and truncation attempts.That created avoidable CPU overhead in a path that executes often, especially with:
This PR removes that repeated invariant work and narrows recomputation to only what changes.
Why this was needed
Behavior and risk
Validation
uv run pytest tests/core/test_context_budget.py tests/agent/test_investigation.py -q --tb=shortmake lintmake format-checkmake typecheck