Skip to content

fix(usage): count dictionary usage in token callbacks - #8171

Merged
Vidit-Ostwal merged 1 commit into
crewAIInc:mainfrom
Winoren:fix/dictionary-token-callback-usage
Oct 23, 2026
Merged

Vidit-Ostwal merged 1 commit into
crewAIInc:mainfrom
Winoren:fix/dictionary-token-callback-usage

Conversation

@Winoren

@Winoren Winoren commented Oct 15, 2026 •

Copy link
Copy Markdown
Contributor

Related issue

Fixes #8170.

Summary

TokenCalcHandler increments successful_requests for dictionary usage but only reads token counters as attributes, leaving prompt, completion, total and cached-prompt counts at zero. This affects the callback accumulator; BaseLLM's separate internal tracker already supports dictionaries.

Reuse UsageMetrics.from_provider_dict for dictionary payloads and accumulate into the existing TokenProcess counters. Keep SDK-object handling, missing usage and absent-process behavior unchanged.

Tests compare dictionary usage with real OpenAI CompletionUsage objects, with and without cached tokens, and verify that repeated callbacks accumulate exactly once per request. They run without provider calls.

Verification

  • Tests added for the changed behavior.
  • Relevant tests and source quality checks pass locally.

Python 3.13.12 / macOS arm64, repository-lockfile dependencies:

# New tests with production unchanged: 2 failed, 5 passed.
uv run --frozen --package crewai pytest lib/crewai/tests/utilities/test_token_counter_callback.py -q

# After fix: 32 passed, 23 existing deprecation warnings.
uv run --frozen --package crewai pytest lib/crewai/tests/utilities/test_token_counter_callback.py lib/crewai/tests/test_usage_shape_parity.py lib/crewai/tests/test_flow_usage_metrics.py -q

uv run --frozen --package crewai ruff check lib/crewai/src/crewai/utilities/token_counter_callback.py
uv run --frozen --package crewai ruff format --check lib/crewai/src/crewai/utilities/token_counter_callback.py
uv run --frozen --package crewai mypy lib/crewai/src/crewai/utilities/token_counter_callback.py
git diff --check

Tests used CREWAI_TELEMETRY_DISABLED=true and OTEL_SDK_DISABLED=true. Full monorepo tests and live LiteLLM integrations were not run. No paid model or provider/registry calls.

Additional context

AI assistance: Codex investigated the code, prepared the implementation and regression tests, and ran the checks above.


Note

Low Risk
Narrow fix to token accounting for dict usage in callbacks; SDK-object and edge-case paths are preserved and covered by new tests.

Overview
Fixes a bug where TokenCalcHandler incremented successful request counts for dict-shaped LLM usage payloads but never updated prompt, completion, or cached token totals because it only read attribute-style SDK objects.

When usage is a dict, the handler now parses it via UsageMetrics.from_provider_dict (same path as BaseLLM) and feeds the resulting metrics into TokenProcess, then returns so object-style handling is unchanged. Missing/empty usage and a nil token process behave as before.

Adds offline tests that dict and CompletionUsage usage accumulate identically (including cached tokens), skip empty usage, and no-op without a TokenProcess.

Reviewed by Cursor Bugbot for commit e5d2a34. Bugbot is set up for automated code reviews on this repo. Configure here.

@strix-security

strix-security Bot commented Oct 15, 2026 •

Copy link
Copy Markdown

Strix Security Review

Warning

This pull request has 55 commits after the last Strix review (b1cf44d). Strix has not reviewed these changes.
Automatic review on push is off for this repository. To review the latest changes, tag @strix-security in a comment, or turn on re-review on push.

No security issues found.

Review summary

Reviewed the PR's two changed files. The production change in lib/crewai/src/crewai/utilities/token_counter_callback.py adds dictionary-usage handling to the token callback by reusing UsageMetrics.from_provider_dict and accumulating prompt, completion, and cached-prompt token counts into the existing TokenProcess counters. The data flow originates from LLM provider usage responses, flows through safe integer coercions (_coerce_int, _first_int), and lands only in in-memory counters — no dangerous sinks and no attacker-controlled input are involved. The added test file is offline counting coverage with no production security pattern. No security vulnerabilities were found.

Updated for b1cf44d.


Reviewed by Strix
Re-run review · Configure security review settings

@coderabbitai

coderabbitai Bot commented Oct 15, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: Organization UI
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 58ee8749-a406-4c42-ac38-f2e6902c0cb6
📥 Commits

Reviewing files that changed from the base of the PR and between b1cf44d and f20d4a0.

📒 Files selected for processing: 1
  • lib/crewai/src/crewai/utilities/token_counter_callback.py

Included review availability: This review used your included allowance. Your plan provides up to 10 included reviews per hour; 6 remain after this review.

📝 Walkthrough

TokenCalcHandler now converts dictionary usage with UsageMetrics and adds prompt, completion, and cached-prompt token counts to TokenProcess. Tests cover dictionary and SDK-object usage, missing usage, and calls without a TokenProcess.

Changes: Token Counting

Layer / File(s) Summary
Usage parsing and accumulation
lib/crewai/src/crewai/utilities/token_counter_callback.py, lib/crewai/tests/utilities/test_token_counter_callback.py
log_success_event converts dictionary usage to UsageMetrics and adds its token counts to TokenProcess. Tests check repeated dictionary and SDK-object events, missing usage, and calls without a TokenProcess.

Priority: ➖ Normal

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to f20d4

Dictionary-shaped usage is accumulated through the same prompt and completion methods as SDK usage, including its total. No merge-blocking risk is established; the PR is ready subject to normal checks.

Pre-merge checks: Passed 4 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 4 passed
Check name Status Explanation
Description check Passed The description includes the related issue, a clear summary of the problem and solution, verification details, test results, scope limitations, and additional context.
Title check Passed The title clearly and concisely identifies the primary change: counting dictionary-shaped usage in token callbacks.
Linked Issues check Passed Issue #8170 requires dictionary usage to record prompt, completion, total, cached-prompt, and successful-request counts. TokenCalcHandler converts dictionary usage with `UsageMetrics.from_provider_d…
Out of Scope Changes check Passed The production change updates dictionary handling in TokenCalcHandler. The added tests cover the linked issue and preserve existing behavior. No unrelated production behavior or unrelated file chang…
  • Fix all pre-merge checks with AI
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create a new PR
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

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

I ran this on Windows with Python 3.13.3 (crewai 1.15.28 from the repo, openai SDK installed, litellm not installed), against main (40bb2d3) and the PR head.

I called TokenCalcHandler.log_success_event({}, {"usage": usage}, 0, 0) with a fresh TokenProcess and printed the counters.

usage passed main PR
OpenAI style dict (12 prompt, 3 completion, 4 cached) requests 1, tokens 0 prompt 12, completion 3, total 15, cached 4
the same data as an openai CompletionUsage object 12, 3, 15, 4 12, 3, 15, 4
Anthropic style dict (input_tokens 10, output_tokens 5, cache_read_input_tokens 20, cache_creation_input_tokens 2) tokens 0 prompt 32, completion 5, total 37, cached 20
Gemini style dict (prompt_token_count 7, candidates_token_count 2) tokens 0 prompt 7, completion 2, total 9
prompt_tokens_details set to None tokens 0 prompt 5, completion 1, cached 0, no error
dict with None token values tokens 0 tokens 0, no error

So the bug is real. On main a dict counts the request but no tokens. The object path is unchanged on the PR.

The new tests/utilities/test_token_counter_callback.py has 7 tests. All 7 pass on the PR. With the file copied onto main, the two dict cases (cached_tokens 0 and 4) fail with assert 0 == 24, and 5 pass.

One thing to know when reading the Anthropic row. from_provider_dict counts cache read and cache creation tokens inside the prompt total (10 + 20 + 2 = 32). So the callback now gives billed prompt tokens for those dicts. The PR says it reuses that helper on purpose, and the number agrees with it. I did not compare it to the BaseLLM tracker at runtime.

I did not run a real LLM call. I also did not check whether the streaming path in llm.py can call the callback and also update the same TokenProcess somewhere else, which would double count. That is read-only and unchecked.

The fix is small and verified for the cases above.

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

Reviewed the dictionary-usage normalization path and its regression coverage. handles nested and provider aliases; accumulation correctly relies on to derive totals. SDK-object behavior remains intact. Approved.

@Vidit-Ostwal
Vidit-Ostwal force-pushed the fix/dictionary-token-callback-usage branch from f20d4a0 to e5d2a34 Compare October 23, 2026 14:04
@Vidit-Ostwal
Vidit-Ostwal merged commit 2f113c0 into crewAIInc:main Oct 23, 2026
59 checks passed
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.

[BUG] TokenCalcHandler drops token counts when usage is a dictionary

3 participants