Skip to content

fix(llmloop): account for tool_calls and arguments in messageTokens context budgeting - #1526

Closed
tokove wants to merge 1 commit into
alibaba:mainfrom
tokove:fix/llmloop-message-tokens-tool-calls
Closed

tokove wants to merge 1 commit into
alibaba:mainfrom
tokove:fix/llmloop-message-tokens-tool-calls

Conversation

@tokove

@tokove tokove commented Sep 22, 2026

Copy link
Copy Markdown

Fixes #1411

AI Disclosure

This PR was developed with assistance from Google DeepMind's Antigravity coding assistant using Gemini 2.5 Pro. The implementation and tests were reviewed manually, and the changes were verified with the project's test suite.

Background & Root Cause

In internal/llmloop/compression.go, context budgeting (specifically messageTokens(m) in CountMessagesTokens and computeActiveZoneSize) estimated token usage using:

return llm.CountTokens(m.ExtractText()) + m.Native.EstimatedTokens()

However:

  1. m.ExtractText() only reads m.Content. For assistant turns containing tool calls, m.Content is typically empty or contains only brief text.
  2. Tool invocation details—including the function name, JSON arguments, and tool call ID—are stored in m.ToolCalls.
  3. Previously, NativeTurn.EstimatedTokens() only accounted for tool use for Anthropic messages when thinking was enabled and for OpenAI responses. Providers using OpenAI chat completions, or Anthropic without thinking, could therefore have tool calls omitted from messageTokens.
  4. As a result, assistant messages with large tool call arguments (for example, detailed code_comment invocations) could be significantly under-counted. This delayed context compression and could ultimately cause upstream LLM requests to fail with HTTP 400 Context Length Exceeded.

Changes

  1. Universal tool call token budgeting: Added toolCallsTokens(calls []llm.ToolCall) int in internal/llmloop/compression.go to count tokens from function names, JSON arguments, and call IDs. Updated messageTokens to include toolCallsTokens(m.ToolCalls).

  2. Eliminated double counting in NativeTurn: Updated NativeTurn.EstimatedTokens() in internal/llm/client.go to estimate only opaque thinking/reasoning blocks. Tool-use estimation was removed from NativeTurn because m.ToolCalls is now handled uniformly across providers.

  3. Unit tests:

    • Added TestCountMessagesTokens_IncludesToolCallsAndArguments and TestComputeActiveZoneSize_AccountsForToolCalls in internal/llmloop/compression_test.go.
    • Added TestNativeTurn_EstimatedTokens in internal/llm/native_turn_test.go.
    • Verified statement coverage above the 90% repository threshold (internal/llm: 95.7%, internal/llmloop: 94.7%).

@github-actions

Copy link
Copy Markdown
Contributor

OpenCodeReview: Review complete: 0 finding(s) across 2 selected item(s).

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

This appears to substantially overlap with #1413, which is already open for the same issue (#1411) and addresses the same messageTokens under-counting of ToolCalls, including the active-zone regression coverage.

The implementation strategy is somewhat different: #1413 centralizes the extra accounting in Message.EstimatedTokens() / ToolCall.EstimatedTokens() and also applies it to the agent/session fallback estimates, while this PR counts ToolCalls directly in llmloop and removes tool-call accounting from NativeTurn.

Before reviewing these as two independent fixes, could you clarify what gap in #1413 this PR is intended to address, or why this implementation should replace that one?

@tokove tokove closed this Sep 22, 2026
@tokove tokove reopened this Sep 22, 2026
@tokove

tokove commented Sep 22, 2026

Copy link
Copy Markdown
Author

Thanks for pointing that out @Qiyuanqiii! I hadn't realized #1413 was already in flight for #1411. Closing this in favor of #1413 to respect the earlier PR.

@tokove tokove closed this Sep 22, 2026
@tokove
tokove deleted the fix/llmloop-message-tokens-tool-calls branch September 22, 2026 02:15
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.

fix(llmloop): account for tool_calls and arguments in messageTokens context budgeting

2 participants