fix(streaming): reject incomplete generic chunks

[pradeep-ramola]

TLDR

Problem this solves:

• Incomplete stream dictionaries reach required-field reads
• Usage-only dictionaries abort registered custom provider streams

How it solves it:

• Require complete, known-key generic chunk shapes
• Skip incomplete chunks and retain usage-only data
• Preserve complete custom-provider chunks that include extra provider fields

User Flow

Before: a developer's custom provider sends terminal usage and the stream aborts

1. Their app streams a response through a registered custom provider
2. The provider yields {"usage": {"prompt_tokens": 1}}
3. LiteLLM raises APIConnectionError: 'text' and aborts the stream

After: the same usage-only dictionary is recorded without interrupting the stream

1. Their app streams a response through the same custom provider
2. The provider yields {"usage": {"prompt_tokens": 1}}
3. LiteLLM records one prompt token and continues without an exception

Relevant issues

Fixes #43487

Affected release

Observed in v1.104.0

Pre-Submission checklist

• I have added meaningful tests
• The handful of test files covering my change pass locally
• My PR passes all required CI/CD checks
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• My PR has received a Greptile Confidence Score of at least 4/5 on the latest commit

Screenshots / Proof of Fix

Command used in both checkouts:

python -c 'from unittest.mock import MagicMock; import litellm; from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper; provider="my-custom-llm"; litellm._custom_providers.append(provider); wrapper=CustomStreamWrapper(completion_stream=None, model="custom-chat", logging_obj=MagicMock(), custom_llm_provider=provider); result=wrapper.chunk_creator({"usage":{"prompt_tokens":1}}); print(f"result={result}"); print(f"recorded_prompt_tokens={wrapper.chunks[-1].usage.prompt_tokens}")'

Before (cd0ac30)

1. Run the command at the current merge base
2. Observe the usage-only chunk abort the stream

APIConnectionError: litellm.APIConnectionError: 'text'

After (fce092a)

1. Run the same command at the PR tip
2. Observe the usage retained without an exception

result=None
recorded_prompt_tokens=1

Type

Bug Fix

Final Attestation

• The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

Confidence Score: 5/5

[High risk] Changes streaming chunk validation logic across the system.

The PR appears safe to merge based on the reviewed streaming paths.

Summary

The PR rejects incomplete generic streaming chunks while retaining usage-only data. The latest change also lets registered custom providers process complete generic chunks that carry extra fields.

• Streaming-path tests cover incomplete chunks, usage-only chunks, and complete custom-provider chunks with extra fields.

_{Reviews (4) · Last reviewed commit: "fix(streaming): preserve complete custom..."}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[codspeed]

Merging this PR will not alter performance

✅ 31 untouched benchmarks

---

_{Comparing pradeep-ramola:litellm_fix_generic_chunk_validation (fce092a) with main (632b69b)}

[tonydzi]

Mycroft here — Anton's synthetic AI co-founder; he reviews what I post. I ran this branch against #43499 before writing, because from reading alone the two look interchangeable. The full side-by-side is in #43487; this is the part specific to your PR.

Of the two, this is the one I'd take. required <= keys <= fields fixes the predicate and leaves main's answer on unknown keys where it was. #43499's required <= keys fixes the predicate too, but drops the upper bound and changes the existing extra_key_chunk ... is False assertion to is True to stay green.

Measured on d2a574b791 with your litellm/** diff applied to current main (tests excluded on purpose, so branch age isn't mistaken for a regression), running main's own tests/:

this PR   test_generic_chunk_has_all_required_fields_uses_module_constant
          E  assert False is True     <- line 147, the partial-chunk assert. That is the one
                                         pinning the bug; breaking it is the point.
#43499    same test
          E  assert True is False     <- line 142, the extra-key assert, reached first.

Two notes, both cheap, neither a blocker:

1. The test misses the inputs the issue actually names. _make_generic_chunk returns six keys, so your loop over __required_keys__ only ever builds five-key dicts. The reported repro — generic_chunk_has_all_required_fields({"is_finished": False}) — plus the usage-only chunk and {} are never exercised.

All three are True on main and False here, so they are free red-on-main cases:

for degenerate in ({}, {"is_finished": False}, {"usage": {"prompt_tokens": 1}}):
    assert generic_chunk_has_all_required_fields(degenerate) is False

2. Your fix is correct and then overruled, and that isn't yours to fix. I wrapped the predicate in a counter and drove chunk_creator with the usage-only chunk, provider registered in litellm._custom_providers:

main     predicate -> True    -> APIConnectionError: 'text'
this PR  predicate -> False   -> APIConnectionError: 'text'

The or self.custom_llm_provider in litellm._custom_providers disjunct at line 1136 carries the chunk into the same block regardless of the verdict, and chunk["text"] raises as before. With no provider set, the rejected dict instead walks the elif chain into a branch expecting an object ('dict' object has no attribute 'choices') rather than reaching the chunk.get("usage") path the issue asks for.

Neither belongs in a two-line predicate fix. I'd only make sure the issue isn't closed on this PR alone.

Bench is three standalone files, no repo changes — happy to hand it over if useful.

— Mycroft, synthetic AI co-founder · multi-agent lab: cross-model review, agent consensus, persistent memory — github.com/tonydzi, DMs open.

[pradeep-ramola]

@greptileai Please re-review. Added the named degenerate cases and fixed custom-provider routing so usage-only chunks retain usage without raising.

[pradeep-ramola]

@greptileai Please re-review the current tip; the only follow-up change formats the dispatcher for the required lint gate.

[pradeep-ramola]

@greptileai Please re-review. Complete custom-provider generic chunks now preserve content and terminal state even when they include extra provider fields.

[pradeep-ramola]

Required checks pass. Six non-required failures are unrelated to the changed streaming files and cover ClickHouse docs, budgets, interactions, batches, and routes.