Skip to content

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

Open
AyushDubey23 wants to merge 1 commit into
alibaba:mainfrom
AyushDubey23:fix/llmloop-token-accounting-tool-calls
Open

AyushDubey23 wants to merge 1 commit into
alibaba:mainfrom
AyushDubey23:fix/llmloop-token-accounting-tool-calls

Conversation

@AyushDubey23

Copy link
Copy Markdown
Contributor

Description

In internal/llmloop/compression.go, messageTokens(m llm.Message) calculates the token weight of messages when partitioning conversation rounds between the active budget and history compression.

func messageTokens(m llm.Message) int {
	return llm.CountTokens(m.ExtractText()) + m.Native.EstimatedTokens()
}

When an assistant turn issues tool calls:

  1. m.ExtractText() returns "" because tool calls and arguments are not stored in m.Content.
  2. For OpenAI-compatible chat completions, m.Native.Payload is either nil or ReasoningPayload (reasoning_content) — it does not store tool calls.
  3. The tool call names, arguments (which often span thousands of characters of file diffs, code writes, or search payloads), and extra metadata live on m.ToolCalls []llm.ToolCall.

Because m.ToolCalls was not accounted for, every assistant tool-call turn evaluated to 0 tokens. In multi-turn reviews with heavy tool usage, computeActiveZoneSize severely under-calculated active zone token usage and packed too many rounds into the context budget, risking context window exhaustion (400 context_length_exceeded).

What Changed

  1. ToolCall.EstimatedTokens(): Added to internal/llm/client.go to estimate tokens for a tool call's name, arguments, and ID framing using the project's standard bytes / 4 heuristic.
  2. Message.EstimatedTokens(): Combines m.Native.EstimatedTokens() with m.ToolCalls token estimation. Added NativeTurn.countsToolCalls() to prevent double-counting when Native.Payload already accounts for tool calls (such as Anthropic MessageParam tool-use blocks or Responses API function-call items).
  3. llmloop.messageTokens(): Updated to use m.EstimatedTokens(), properly accounting for tool calls and arguments in active zone size calculation (computeActiveZoneSize) and CountMessagesTokens.
  4. Session history & agent consistency: Aligned fallback token estimation in internal/session/history.go and internal/agent/util.go to include m.EstimatedTokens().

Tests Added

  • internal/llm/client_test.go:
    • TestToolCall_EstimatedTokens: Covers empty, short, and large tool-call argument payloads.
    • TestMessage_EstimatedTokens: Verifies text-only, tool-call, mixed reasoning+tool-call messages, and guards against double-counting with Anthropic native tool use blocks.
  • internal/llmloop/compression_test.go:
    • TestCountMessagesTokens_IncludesToolCalls: Asserts that assistant messages with tool calls produce accurate, non-zero token counts.
    • TestComputeActiveZoneSize_AccountsForToolCalls: Verifies that rounds with large tool-call arguments consume active-zone budget properly rather than being treated as 0 tokens.

Checklist

  • My code follows the project's coding style (gofmt, go vet)
  • All unit tests pass locally (go test ./internal/...)
  • Line endings are normalized to LF
  • verify-english-only.go passes cleanly
  • I did not attribute commits to AI/LLM, and I understand and can explain all code written

Closes #1411

Assistant turns issuing tool calls under OpenAI-compatible chat completions store their calls and arguments on ToolCalls rather than Native.Payload. Because messageTokens only inspected ExtractText() and Native.EstimatedTokens(), these turns evaluated to 0 tokens during active zone budget calculations, causing computeActiveZoneSize to pack too many rounds and risk exceeding model context limits.

Adds EstimatedTokens to ToolCall and Message, and updates messageTokens to account for tool calls when not already covered by Native payload.
@github-actions

Copy link
Copy Markdown
Contributor

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

@NanaseInori NanaseInori left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

This addresses the missing tool-call contribution to context budgeting. The native-payload guard avoids counting tool calls again when Anthropic or Responses payloads already include them, and the agent/session fallback estimates use the same accounting path.

The regression tests cover large arguments consuming the active-zone budget instead of being treated as zero. No blocking findings from me; approved.

This review was conducted by Qiyuanqiii's review bot, using the model GPT 6 Astra Max. If you need a human review, please manually @.

@AyushDubey23

Copy link
Copy Markdown
Contributor Author

Hi @Qiyuanqiii, the bot review has passed and all 13 CI checks are green. Could you please take a quick look and provide the human write-access approval/merge when you have a moment? Thanks!

@Qiyuanqiii

Copy link
Copy Markdown
Contributor

Waiting for the maintainer themselves to review it; he may not have time to look at it.

This branch has not been deployed

No deployments
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

3 participants