[Pranav Tomar (Mystery01092000)]

Related issue

Fixes #9186

Summary

Azure usage lost cached_prompt_tokens and reasoning_tokens in two ways:

• Streaming (sync and async): _handle_streaming_completion and _ahandle_streaming_completion built the usage dict from prompt_tokens, completion_tokens and total_tokens only, so cached and reasoning counts were always 0. Both now call _extract_azure_token_usage on the usage chunk, the same extraction the non-streaming path uses.
• Non-streaming as well: _extract_azure_token_usage read usage.prompt_tokens_details and usage.completion_tokens_details as attributes. azure-ai-inference's CompletionsUsage (1.0.0b9) does not declare those fields, so on a real SDK object they exist only as mapping items and the attribute lookup returns None. The existing test passed because it builds the response with MagicMock. A new _usage_detail helper reads the detail as an attribute or a mapping item, so both mocked and real objects work.

Verification

• Tests added or updated for the changed behavior
  • lib/crewai/tests/llms/azure/test_azure_streaming_usage.py: a sync and an async streaming call fed real StreamingChatCompletionsUpdate chunks (content, then a usage chunk with prompt_tokens_details.cached_tokens=40 and completion_tokens_details.reasoning_tokens=60) record both counts in get_token_usage_summary(); and _extract_azure_token_usage reads both from a real ChatCompletions object. All three fail on main and pass with this change.
• Relevant tests and quality checks pass locally
  • uv run pytest lib/crewai/tests/llms/azure --allowed-hosts=127.0.0.1: 91 passed
  • uv run pytest lib/crewai/tests/llms --allowed-hosts=127.0.0.1: 1068 passed, 20 skipped, 2 failed. The failures are in test_missing_api_key.py and the Anthropic/Bedrock PDF multimodal tests; they fail the same way with the unmodified completion.py and do not touch Azure.
  • uv run ruff check / uv run ruff format --check on the changed files: clean
  • uv run mypy --config-file pyproject.toml lib/crewai/src/crewai/llms/providers/azure/completion.py: no issues

Additional context

• Tested on Windows 11. The async tests need --allowed-hosts=127.0.0.1 locally, because the Windows Proactor event loop opens a loopback socketpair that --block-network rejects. Not needed on Linux CI.
• AI assistance was used for this change (please apply llm-generated; I can't set labels on this repo).

🤖 Generated with Claude Code

---

Note

Low Risk
Changes are limited to Azure usage extraction and metrics; no auth, request routing, or response content logic is affected.

Overview
Azure token usage reporting now preserves cached_prompt_tokens and reasoning_tokens for streaming calls and when the Inference SDK exposes detail fields only as dict keys.

Sync and async streaming handlers no longer build usage from the three totals alone; they route usage chunks through _extract_azure_token_usage, matching the non-streaming path. A new _usage_detail helper reads nested prompt_tokens_details / completion_tokens_details via attributes or Mapping lookup, because real CompletionsUsage objects omit those fields as attributes. Tests cover sync/async streaming and non-streaming extraction with SDK-shaped usage payloads.

^{Reviewed by Cursor Bugbot for commit 38514c1. Bugbot is set up for automated code reviews on this repo. Configure here.}

[strix-security]

Strix Security Review

Warning

This pull request has 27 commits after the last Strix review (8a5ff09). 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 open security findings.

PR summary: Fixes Azure token usage reporting so cached-prompt and reasoning token counts are preserved in both streaming and non-streaming paths, adding a `_usage_detail` helper that reads nested usage details as attributes or mapping items, plus tests using real `azure-ai-inference` objects.

Review summary

This PR fixes Azure token accounting so cached-prompt and reasoning token counts survive streaming: both streaming handlers now call _extract_azure_token_usage on the usage chunk, and a new _usage_detail helper reads prompt_tokens_details / completion_tokens_details as either attributes or mapping items (the SDK's CompletionsUsage only exposes them as mapping items).

I read the full diff and traced the new data flow end to end: usage chunk → _extract_azure_token_usage → _finalize_streaming_response → _track_token_usage_internal → UsageMetrics.from_provider_dict (_coerce_int / _first_int, which collapse malformed values to 0 and recompute total_tokens from the prompt + completion counts) → LLMCallCompletedEvent.usage and telemetry span attributes. Downstream consumers read every bucket with .get() or 0, no consumer depends on the old three-key shape, and no security control (rate limiting is request-count based, and there is no token budget or quota enforcement) consumes these counters. The widened signature changes nothing at the other call site, and its only input is the Azure inference response over TLS with API-key/Entra auth — not end-user input — so nothing here is attacker-reachable.

The lint triage returned no item for any code this PR adds or changes; the single item in the changed file is 1,100 lines away in _normalize_azure_fields (endpoint/API-key handling the changed functions never call), and the remaining items (BaseLLM.api_key serialization, unclamped Retry-After, fail-open hook dispatch) are pre-existing backlog reachable from any provider call and untouched here. No security issues found in the changed code.

Ruled out (1)

• Can malformed or hostile provider usage data (non-numeric or mapping values) from the new `_usage_detail` path reach an unsafe sink or weaken a token-based con…
  Every count is routed through `UsageMetrics.from_provider_dict`'s `_coerce_int`/`_first_int` coercers, consumers read buckets with `.get() or 0`, and no token budget or quota enforcement exists — only reporting and telemetry, with the input being the authenticated Azure inference response rather th…

_{Review 1 on this PR · 8a5ff09 · requested Nov 7 12:47 UTC when the PR was opened}

---

Reviewed by Strix
_{Re-run review · Configure security review settings}

[Christopher Krah (chrikrah)]

Pranav Tomar (@Mystery01092000) nothing here blocks a merge from my side at 8a5ff09. With completion.py reverted and your tests kept, all three new tests fail:

# uv sync --frozen --package crewai --extra azure-ai-inference --dev, azure-ai-inference 1.0.0b9, Python 3.12
$ python -m pytest lib/crewai/tests/llms/azure -q -p no:cacheprovider -n 4
# head 8a5ff09
118 passed, 102 warnings in 14.55s
# head with completion.py reverted to 73f4832, tests kept
FAILED test_azure_streaming_usage.py::test_non_streaming_usage_reads_details_from_the_sdk_model
FAILED test_azure_streaming_usage.py::test_async_streaming_usage_keeps_cached_and_reasoning_tokens
FAILED test_azure_streaming_usage.py::test_streaming_usage_keeps_cached_and_reasoning_tokens
3 failed, 115 passed, 102 warnings in 14.26s

# probe at head: real ChatCompletions, usage with no *_details, then with prompt_tokens_details=None
$ python -I probe9211.py
{'prompt_tokens': 5, 'completion_tokens': 7, 'total_tokens': 12, 'cached_prompt_tokens': 0, 'reasoning_tokens': 0}
{'prompt_tokens': 5, 'completion_tokens': 7, 'total_tokens': 12, 'cached_prompt_tokens': 0, 'reasoning_tokens': 0}
# not run: a live Azure deployment

The detail reads came with 0341e5a (#4447) and fc6792d (#5389), and their tests build usage as a MagicMock, which answers any attribute. Your non-streaming test is the first against a real ChatCompletions.

non-blocking: test_azure_no_detail_fields pins the no-details case only on a MagicMock. If you want it on the real model, the two probe objects drop into the new test file unchanged.

[Pranav Tomar (Mystery01092000)]

Hi João Moura (@joaomdmoura), could you review this one? It keeps cached-prompt and reasoning token counts in Azure usage, for both streaming and non-streaming. Christopher Krah (@chrikrah) reproduced it above: with completion.py reverted, all three new tests fail, and they pass with the fix.