fix(core): prevent inflated token counts when merging streaming usage_metadata by ShresthSamyak · Pull Request #35615 · langchain-ai/langchain

Shresth Samyak (ShresthSamyak) · 2026-03-07T00:42:18Z

Fix incorrect token counting when streaming with stream_usage=True.

Root cause:
Streaming providers (e.g. OpenAI) return cumulative token usage in each chunk.
add_ai_message_chunks() previously summed usage_metadata across chunks,
which inflated token counts.

Example:
chunk1: input_tokens=10
chunk2: input_tokens=10
chunk3: input_tokens=10

Previous result:
30 input tokens

Correct result:
10 input tokens (last cumulative value)

Fix:
Use the last non-null usage_metadata when merging message chunks.

Adds a regression test ensuring streaming token usage is not summed.

Fixes #31351

…_metadata

Copilot

Pull request overview

Fixes incorrect (inflated) token usage totals when merging AIMessageChunks produced during streaming with stream_usage=True, where providers emit cumulative usage_metadata in each chunk.

Changes:

Update add_ai_message_chunks() to merge usage_metadata by selecting the last non-null value instead of summing across chunks.
Add a regression test suite covering cumulative streaming usage, all-None usage, and partial usage across chunks.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
`libs/core/langchain_core/messages/ai.py`	Changes `usage_metadata` merge behavior in `add_ai_message_chunks()` to avoid inflating token totals for cumulative streaming usage.
`libs/core/tests/unit_tests/messages/test_usage_metadata_streaming.py`	Adds regression tests to ensure streaming usage is not summed and that None/partial usage cases behave as intended.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-03-07T00:45:43Z

+    # Take the last non-null usage_metadata rather than summing, because
+    # streaming providers (e.g. OpenAI with stream_usage=True) send
+    # cumulative token counts in each chunk. Summing them would inflate
+    # the totals. See: https://github.com/langchain-ai/langchain/issues/31351
+    usage_metadata: UsageMetadata | None = None
+    for candidate in reversed((left, *others)):
+        if candidate.usage_metadata is not None:
+            usage_metadata = candidate.usage_metadata
+            break



add_ai_message_chunks used to aggregate usage_metadata across chunks via add_usage, but this change now always selects the last non-null value. That is a behavior change that conflicts with existing unit expectations (e.g. libs/core/tests/unit_tests/messages/test_ai.py::test_add_ai_message_chunks_usage currently expects summation) and may break callers/providers that emit per-chunk delta usage. Please either (a) update the existing tests + document this as an intentional behavior change, or (b) keep backwards compatibility by supporting both cumulative and delta strategies (e.g., gate on explicit metadata/flag, or detect and fall back to summation when usage values look like deltas).

Suggested change

# Take the last non-null usage_metadata rather than summing, because

# streaming providers (e.g. OpenAI with stream_usage=True) send

# cumulative token counts in each chunk. Summing them would inflate

# the totals. See: https://github.com/langchain-ai/langchain/issues/31351

usage_metadata: UsageMetadata | None = None

for candidate in reversed((left, *others)):

if candidate.usage_metadata is not None:

usage_metadata = candidate.usage_metadata

break

# Support both cumulative and per-chunk (delta) usage aggregation.

# Some streaming providers (e.g. OpenAI with stream_usage=True) send

# cumulative token counts in each chunk, while others may emit per-chunk

# deltas. To preserve backwards compatibility, we:

# - Detect sequences that look cumulative (monotonically non-decreasing

# per key) and take the last such value.

# - Otherwise, fall back to summing per-chunk values.

def _looks_cumulative(usages: list[UsageMetadata]) -> bool:

"""Heuristically determine whether a sequence of usage dicts is cumulative.

A sequence is considered cumulative if, for every key, values never

decrease across the sequence and at least one key increases at least once.

"""

if not usages:

return False

last_values: dict[str, int] = {}

seen_increase = False

for usage in usages:

for key, value in usage.items():

prev = last_values.get(key, 0)

if value < prev:

# Decrease detected -> not cumulative.

return False

if value > prev:

seen_increase = True

last_values[key] = value

return seen_increase

usage_metadata: UsageMetadata | None = None

usage_candidates = [

msg.usage_metadata

for msg in (left, *others)

if msg.usage_metadata is not None

]

if usage_candidates:

if _looks_cumulative(usage_candidates):

# Cumulative sequence: last value already contains full usage.

usage_metadata = usage_candidates[-1]

else:

# Delta sequence: sum values per key (backwards compatible behavior).

merged = usage_candidates[0]

for u in usage_candidates[1:]:

merged = cast(

"UsageMetadata", _dict_int_op(merged, u, operator.add)

)

usage_metadata = merged

Copilot · 2026-03-07T00:45:44Z

+    usage_metadata: UsageMetadata | None = None
+    for candidate in reversed((left, *others)):
+        if candidate.usage_metadata is not None:
+            usage_metadata = candidate.usage_metadata
+            break


usage_metadata is a mutable dict, and this code assigns it directly from one of the input chunks. That means the merged chunk and the original chunk will share the same dict object, so later mutations to either (even accidental) will affect the other. Consider copying (and copying nested detail dicts as needed) when setting usage_metadata so the merged result is independent, similar to how the previous add_usage path always produced a new object.

Copilot · 2026-03-07T00:45:44Z

+def test_usage_metadata_streaming_not_summed() -> None:
+    """Merged usage_metadata should equal the last chunk, not the sum."""
+    chunk1 = AIMessageChunk(
+        content="Hello",
+        usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11},
+    )
+
+    chunk2 = AIMessageChunk(
+        content=" world",
+        usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12},
+    )
+
+    chunk3 = AIMessageChunk(
+        content="!",
+        usage_metadata={"input_tokens": 10, "output_tokens": 3, "total_tokens": 13},
+    )
+
+    merged = add_ai_message_chunks(chunk1, chunk2, chunk3)
+
+    assert merged.usage_metadata is not None
+    assert merged.usage_metadata["input_tokens"] == 10
+    assert merged.usage_metadata["output_tokens"] == 3
+    assert merged.usage_metadata["total_tokens"] == 13


This new regression test asserts the 'last chunk wins' semantics, but the repository already has a unit test that asserts add_ai_message_chunks sums usage across chunks (libs/core/tests/unit_tests/messages/test_ai.py::test_add_ai_message_chunks_usage). With the current code changes, that existing test will fail in CI. Please update/replace the older test (or adjust implementation to support both semantics) so the test suite reflects a single intended contract.

codspeed-hq · 2026-03-07T00:45:46Z

Merging this PR will not alter performance

⚠️

Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

✅ 13 untouched benchmarks
⏩ 23 skipped benchmarks¹

_{Comparing ShresthSamyak:fix-streaming-token-count (ce9e7a3) with master (29134dc)}

23 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

…ut breaking normal merge behavior

Shresth Samyak (ShresthSamyak) · 2026-03-07T04:52:32Z

Eugene Yurtsev (@eyurtsev) all checks clear

Shresth Samyak (ShresthSamyak) · 2026-03-07T19:50:45Z

ccurme (@ccurme) all checks clear

Yu Liu (shell-nlp) · 2026-03-23T02:22:59Z

Hope to merge soon

Kevin Kakolla (torontodeveloper) · 2026-04-12T18:53:32Z

+    # counts (e.g. OpenAI with stream_usage=True sends running totals in
+    # every chunk) vs. independent invocations that should be summed.
+    # See: https://github.com/langchain-ai/langchain/issues/31351
+    usages = [m.usage_metadata for m in (left, *others) if m.usage_metadata is not None]


I know u have not introduced, these single char variable names 'm' should be Discouraged. i have no clue what is 'm' here and tried to look previous lines but it seems it's there some where.
Also just curious instead of m.usage_metadata is not None, is if m.usage_metadata returns truthy is not enough? or equivalent

Kevin Kakolla (torontodeveloper) · 2026-04-12T18:56:10Z

+    # See: https://github.com/langchain-ai/langchain/issues/31351
+    usages = [m.usage_metadata for m in (left, *others) if m.usage_metadata is not None]
+
+    if not usages:


do we get usages is None or probably empty list?. I am not sure when we get 'usages' is empty list. probably if m is empty ?. My concern seems does this if condiiton ever executes? Even then we are reassign usage_metadata to none based on usages which depends on m.usage_metadata. there seems some circling going on which I am not able pin point

Kevin Kakolla (torontodeveloper) · 2026-04-12T18:56:56Z

+        inputs = [u.get("input_tokens", 0) for u in usages]
+        totals = [u.get("total_tokens", 0) for u in usages]


these single char variable names is not Best practice

Kevin Kakolla (torontodeveloper) · 2026-04-12T18:59:25Z

+        else:
+            usage_metadata = usages[0].copy()
+            for u in usages[1:]:
+                usage_metadata = add_usage(usage_metadata, u)


what does usages list contains? is it usage of token or metadata? I would use that name as variable than single character names.

I would put note or TO-DO for future refactoring this file is way too long with almost 900 lines, hard to maintain/readable

Kevin Kakolla (torontodeveloper) · 2026-04-12T19:00:14Z

+    assert merged.usage_metadata is not None
+    assert merged.usage_metadata["input_tokens"] == 5
+    assert merged.usage_metadata["output_tokens"] == 7
+    assert merged.usage_metadata["total_tokens"] == 12


just curious what are these numbers meant? 5,7,12 etc.

Kevin Kakolla (torontodeveloper)

.

fix(core): prevent inflated token counts when merging streaming usage…

26a849c

…_metadata

Copilot AI review requested due to automatic review settings March 7, 2026 00:42

Shresth Samyak (ShresthSamyak) requested a review from Eugene Yurtsev (eyurtsev) as a code owner March 7, 2026 00:42

github-actions Bot added core `langchain-core` package issues & PRs fix For PRs that implement a fix external labels Mar 7, 2026

Copilot started reviewing on behalf of Shresth Samyak (ShresthSamyak) March 7, 2026 00:42 View session

Copilot AI reviewed Mar 7, 2026

View reviewed changes

Shresth Samyak (ShresthSamyak) and others added 4 commits March 7, 2026 06:24

fix(core): handle cumulative usage_metadata in streaming merges witho…

98321ca

…ut breaking normal merge behavior

fix(core): correctly merge cumulative and delta usage_metadata

12ec745

style(core): fix lint formatting

17fb544

Merge branch 'master' into fix-streaming-token-count

ce9e7a3

github-actions Bot added the size: S 50-199 LOC label Mar 9, 2026

Kevin Kakolla (torontodeveloper) reviewed Apr 12, 2026

View reviewed changes

-    # Take the last non-null usage_metadata rather than summing, because
-    # streaming providers (e.g. OpenAI with stream_usage=True) send
-    # cumulative token counts in each chunk. Summing them would inflate
-    # the totals. See: https://github.com/langchain-ai/langchain/issues/31351
-    usage_metadata: UsageMetadata | None = None
-    for candidate in reversed((left, *others)):
-        if candidate.usage_metadata is not None:
-            usage_metadata = candidate.usage_metadata
-            break
+    # Support both cumulative and per-chunk (delta) usage aggregation.
+    # Some streaming providers (e.g. OpenAI with stream_usage=True) send
+    # cumulative token counts in each chunk, while others may emit per-chunk
+    # deltas. To preserve backwards compatibility, we:
+    #   - Detect sequences that look cumulative (monotonically non-decreasing
+    #     per key) and take the last such value.
+    #   - Otherwise, fall back to summing per-chunk values.
+    def _looks_cumulative(usages: list[UsageMetadata]) -> bool:
+        """Heuristically determine whether a sequence of usage dicts is cumulative.
+        A sequence is considered cumulative if, for every key, values never
+        decrease across the sequence and at least one key increases at least once.
+        """
+        if not usages:
+            return False
+        last_values: dict[str, int] = {}
+        seen_increase = False
+        for usage in usages:
+            for key, value in usage.items():
+                prev = last_values.get(key, 0)
+                if value < prev:
+                    # Decrease detected -> not cumulative.
+                    return False
+                if value > prev:
+                    seen_increase = True
+                last_values[key] = value
+        return seen_increase
+    usage_metadata: UsageMetadata | None = None
+    usage_candidates = [
+        msg.usage_metadata
+        for msg in (left, *others)
+        if msg.usage_metadata is not None
+    ]
+    if usage_candidates:
+        if _looks_cumulative(usage_candidates):
+            # Cumulative sequence: last value already contains full usage.
+            usage_metadata = usage_candidates[-1]
+        else:
+            # Delta sequence: sum values per key (backwards compatible behavior).
+            merged = usage_candidates[0]
+            for u in usage_candidates[1:]:
+                merged = cast(
+                    "UsageMetadata", _dict_int_op(merged, u, operator.add)
+                )
+            usage_metadata = merged

		inputs = [u.get("input_tokens", 0) for u in usages]
		totals = [u.get("total_tokens", 0) for u in usages]

Uh oh!

Conversation

Shresth Samyak (ShresthSamyak) commented Mar 7, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Mar 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 7, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Mar 7, 2026

Choose a reason for hiding this comment

Uh oh!

codspeed-hq Bot commented Mar 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will not alter performance

Footnotes

Uh oh!

Shresth Samyak (ShresthSamyak) commented Mar 7, 2026

Uh oh!

Shresth Samyak (ShresthSamyak) commented Mar 7, 2026

Uh oh!

Yu Liu (shell-nlp) commented Mar 23, 2026

Uh oh!

Kevin Kakolla (torontodeveloper) Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Kakolla (torontodeveloper) Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Kakolla (torontodeveloper) Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Kakolla (torontodeveloper) Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Kakolla (torontodeveloper) Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

Kevin Kakolla (torontodeveloper) left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

codspeed-hq Bot commented Mar 7, 2026 •

edited

Loading