PR #26245: fix(proxy): skip redundant tiktoken recount when provider supplies reasoning_tokens

• Repository: BerriAI/litellm
• Author: dschulmeist
• State: open | merged: False
• Link: https://github.com/BerriAI/litellm/pull/26245

Description (problem / solution / changelog)

Relevant issues

Fixes #26193

Type

Bug Fix

Changes

Problem

stream_chunk_builder unconditionally calls ChunkProcessor.count_reasoning_tokens, which recomputes reasoning_tokens via tiktoken.encode() — a C extension that holds the GIL for tens of seconds on large reasoning responses (Claude extended thinking, OpenAI o1/o3, Gemini thinking, 100k+ tokens).

The computed value is already discarded by calculate_usage when the provider supplied reasoning_tokens in a streaming usage chunk. Specifically, calculate_usage only applies the recomputed value when completion_tokens_details.reasoning_tokens is None:

if reasoning_tokens is not None:
    if returned_usage.completion_tokens_details is None:
        returned_usage.completion_tokens_details = (
            CompletionTokensDetailsWrapper(reasoning_tokens=reasoning_tokens)
        )
    elif (
        returned_usage.completion_tokens_details is not None
        and returned_usage.completion_tokens_details.reasoning_tokens is None
    ):
        returned_usage.completion_tokens_details.reasoning_tokens = (
            reasoning_tokens
        )

All major reasoning providers (Anthropic extended thinking, OpenAI o1/o3, Gemini thinking) supply reasoning_tokens in streaming usage chunks, so in the common case the expensive recount runs purely to be thrown away.

The bug report (#26193) includes py-spy thread dumps showing a single tiktoken.encode() call monopolizing the event loop for the full liveness-probe window, resulting in pods being SIGKILLed by Kubernetes.

Fix

Inspect the already-parsed streaming chunks for any completion_tokens_details.reasoning_tokens entry. If one exists, skip the local tiktoken recount entirely — calculate_usage will pick up the provider-supplied value as before. If none exists (rare), fall through to count_reasoning_tokens so providers that emit reasoning_content without reporting usage are unaffected.

New helper ChunkProcessor.chunks_have_reasoning_tokens(chunks) handles both ModelResponseStream objects and plain-dict chunks, matching the existing chunk-parsing patterns in _calculate_usage_per_chunk.

Call-site change in litellm/main.py:

if processor.chunks_have_reasoning_tokens(chunks):
    reasoning_tokens = None
else:
    reasoning_tokens = processor.count_reasoning_tokens(response)

Tests

Four new tests in tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py:

1. test_chunks_have_reasoning_tokens_true_when_provider_supplies — regression test for the reported bug.
2. test_chunks_have_reasoning_tokens_false_when_provider_omits — ensures the slow fallback is still exercised when providers omit usage.
3. test_chunks_have_reasoning_tokens_handles_dict_chunks — dict-chunk path (proxy-style chunks) is covered.
4. test_stream_chunk_builder_skips_count_reasoning_tokens_when_usage_present — end-to-end: stream_chunk_builder does not invoke count_reasoning_tokens when the provider already supplied the value, and the final response carries the provider-supplied number.

All four tests fail on litellm_oss_branch without the fix and pass with it.

Scope

• No API changes: existing public method signatures untouched.
• No new config flags or defaults.
• No changes to count_reasoning_tokens itself (it remains the correct fallback).
• Does not address the separate case where a misbehaving provider emits reasoning_content without any usage chunk. That path still blocks the event loop and could be fixed in a follow-up (e.g. asyncio.to_thread at the async caller), but is orthogonal to this issue's reported reproduction.

Pre-Submission checklist

• Added testing in tests/test_litellm/
• PR passes the touched test file (13/13 in test_streaming_chunk_builder_utils.py)
• Scope isolated to a single bug
• Greptile review (auto-triggered)

Changed files

• litellm/litellm_core_utils/streaming_chunk_builder_utils.py (modified, +55/-9)
• litellm/main.py (modified, +9/-1)
• tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py (modified, +128/-11)