fix(streaming): avoid duplicate usage metadata chunk

[Pouyanpi]

Description

avoid duplicate usage metadata chunk. addressing #1977 (comment)

Related Issue(s)

#1977 (comment)

Verification

AI Assistance

• No AI tools were used.
• AI tools were used; a human reviewed and can explain every change (tool: ___).

Checklist

• I've read the CONTRIBUTING guidelines.
• This PR links to a triaged issue assigned to me.
• My PR title follows the project commit convention.
• I've updated the documentation if applicable.
• I've added tests if applicable.
• I've noted any verification beyond CI and any checks I couldn't run.
• I did not update generated changelog files manually.
• I addressed all CodeRabbit, Greptile, and other review comments, or replied with why no change is needed.
• @mentions of the person or team responsible for reviewing proposed changes.

Summary by CodeRabbit

• Bug Fixes

  • Fixed streaming metadata so usage information is only included once and does not repeat in the final streamed chunk.
  • Improved handling of streamed metadata when metadata display is enabled, preventing stale usage data from carrying over to later chunks.

• Tests

  • Added coverage to confirm raw usage metadata appears a single time and is not duplicated at stream completion.

[greptile-apps]

Greptile Summary

This PR fixes a duplicate-usage-metadata bug in StreamingHandler: when a provider sends usage statistics in a non-final chunk, the data was being re-attached verbatim to the terminal END_OF_STREAM chunk as well. The fix pops "usage" from current_metadata immediately after the non-final chunk is enqueued, so the final chunk only carries response_metadata and usage_metadata defaults.

• nemoguardrails/streaming.py: Two lines added after a non-final metadata chunk is enqueued to remove "usage" from self.current_metadata, preventing duplication in the final END_OF_STREAM chunk.
• tests/test_streaming_handler.py: New async test pushes a usage-bearing empty chunk followed by the stream terminator and asserts that exactly one chunk carries usage and the final chunk does not.

Confidence Score: 5/5

The change is minimal and precisely targeted — two guarded lines in a single branch of _process — with a dedicated new test that directly exercises the fixed scenario.

The mutation to current_metadata happens after the chunk is already enqueued, so no already-sent data is affected. The guard (chunk is not END_OF_STREAM and "usage" in chunk_dict.get("metadata", {})) is tight and avoids interfering with the final-chunk metadata path. The new test confirms both that exactly one usage chunk is emitted and that the final chunk carries only the expected default keys.

No files require special attention; the only noteworthy item is a minor truthiness vs. is not None style nit in the new test.

Important Files Changed

Filename	Overview
nemoguardrails/streaming.py	Adds two lines after enqueuing a non-final metadata chunk to pop usage from current_metadata, preventing it from being re-emitted verbatim in the final END_OF_STREAM metadata chunk.
tests/test_streaming_handler.py	New test test_raw_usage_metadata_chunk_is_not_repeated_on_final_chunk verifies the fix. Assertions correctly check the final chunk via chunks[-1], but the usage_chunks filter relies on truthiness rather than an explicit is not None check.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant P as Provider
    participant SH as StreamingHandler
    participant Q as Queue
    participant C as Consumer

    P->>SH: push_chunk("Hello")
    SH->>SH: "_process("Hello") — current_metadata={}"
    SH->>Q: "enqueue {text:"Hello"} (no metadata)"

    P->>SH: "push_chunk("", metadata={"usage": {...}})"
    SH->>SH: "current_metadata.update({"usage": {...}})"
    SH->>SH: "_process("") — current_metadata={"usage":{...}}"
    SH->>Q: "enqueue {text:"", metadata:{"usage":{...}}}"
    SH->>SH: current_metadata.pop("usage") ← FIX

    P->>SH: push_chunk(None) → END_OF_STREAM
    SH->>SH: "_process(END_OF_STREAM) — current_metadata={}"
    SH->>Q: "enqueue {text:"", metadata:{response_metadata:None, usage_metadata:None}}"

    Q-->>C: "{text:"Hello"}"
    Q-->>C: "{text:"", metadata:{"usage":{...}}}"
    Q-->>C: "{text:"", metadata:{response_metadata:None, usage_metadata:None}}"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant P as Provider
    participant SH as StreamingHandler
    participant Q as Queue
    participant C as Consumer

    P->>SH: push_chunk("Hello")
    SH->>SH: "_process("Hello") — current_metadata={}"
    SH->>Q: "enqueue {text:"Hello"} (no metadata)"

    P->>SH: "push_chunk("", metadata={"usage": {...}})"
    SH->>SH: "current_metadata.update({"usage": {...}})"
    SH->>SH: "_process("") — current_metadata={"usage":{...}}"
    SH->>Q: "enqueue {text:"", metadata:{"usage":{...}}}"
    SH->>SH: current_metadata.pop("usage") ← FIX

    P->>SH: push_chunk(None) → END_OF_STREAM
    SH->>SH: "_process(END_OF_STREAM) — current_metadata={}"
    SH->>Q: "enqueue {text:"", metadata:{response_metadata:None, usage_metadata:None}}"

    Q-->>C: "{text:"Hello"}"
    Q-->>C: "{text:"", metadata:{"usage":{...}}}"
    Q-->>C: "{text:"", metadata:{response_metadata:None, usage_metadata:None}}"

Loading

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
tests/test_streaming_handler.py:747
The filter uses a truthiness check on the `usage` dict. If all token counts were zero, the dict `{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}` would still be truthy (non-empty dict), so the current test data passes — but relying on truthiness here is subtly misleading. An explicit `is not None` guard is more accurate and intention-revealing.

```suggestion
        usage_chunks = [chunk for chunk in chunks if chunk.get("metadata", {}).get("usage") is not None]
```

_{Reviews (1): Last reviewed commit: "fix(streaming): avoid duplicate usage me..." | Re-trigger Greptile}

[greptile-apps]

The filter uses a truthiness check on the usage dict. If all token counts were zero, the dict {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} would still be truthy (non-empty dict), so the current test data passes — but relying on truthiness here is subtly misleading. An explicit is not None guard is more accurate and intention-revealing.

Suggested change

	usage_chunks = [chunk for chunk in chunks if chunk.get("metadata", {}).get("usage")]
	usage_chunks = [chunk for chunk in chunks if chunk.get("metadata", {}).get("usage") is not None]

Prompt To Fix With AI

This is a comment left during a code review.
Path: tests/test_streaming_handler.py
Line: 747

Comment:
The filter uses a truthiness check on the `usage` dict. If all token counts were zero, the dict `{"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}` would still be truthy (non-empty dict), so the current test data passes — but relying on truthiness here is subtly misleading. An explicit `is not None` guard is more accurate and intention-revealing.

```suggestion
        usage_chunks = [chunk for chunk in chunks if chunk.get("metadata", {}).get("usage") is not None]
```

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c94c6577-1ac3-4d5e-aee2-1bf20228b0dd

📥 Commits

Reviewing files that changed from the base of the PR and between 15f87ee and 6043437.

📒 Files selected for processing (2)

• nemoguardrails/streaming.py
• tests/test_streaming_handler.py

---

📝 Walkthrough

Walkthrough

The stream handler now removes raw usage metadata after emitting a non-final chunk. A new async test verifies that this metadata appears once and is not repeated on the final chunk.

Changes

Usage metadata reset

Layer / File(s)	Summary
Clear usage metadata after enqueue nemoguardrails/streaming.py, tests/test_streaming_handler.py	usage is removed from ongoing chunk metadata after a non-final enqueue, and a regression test confirms it is not duplicated onto the final chunk.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: preventing duplicate usage metadata chunks in streaming.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Test Results For Major Changes	✅ Passed	Minor bug fix in streaming metadata handling with a new async test; the description doesn't need separate test results for this scope.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/streaming-usage-metadata-once

---

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