Skip to content

Fix: streaming usage tracking#90

Merged
SantiagoDePolonia merged 3 commits intomainfrom
fix/streaming-usage-tracking
Feb 25, 2026
Merged

Fix: streaming usage tracking#90
SantiagoDePolonia merged 3 commits intomainfrom
fix/streaming-usage-tracking

Conversation

@SantiagoDePolonia
Copy link
Copy Markdown
Contributor

@SantiagoDePolonia SantiagoDePolonia commented Feb 25, 2026

Summary by CodeRabbit

  • Breaking Changes

    • Streaming response finalization event renamed from "response.done" to "response.completed" across providers and tests
  • Improvements

    • More reliable incremental parsing of streaming responses to better preserve usage data
    • Reduced memory impact when handling large or fragmented streaming responses

SantiagoDePolonia and others added 2 commits February 25, 2026 00:59
…ng in StreamUsageWrapper

The rolling buffer approach could truncate large SSE events (like the
Responses API's final event containing the full response object + usage).
Switch to incremental parsing that processes complete events on each
\n\n boundary, caching usage data as it's found.

Also adds usage field to mock provider's streaming response.done event
and new tests for Responses API format, large events, and fragmented reads.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
… in /v1/responses

OpenAI's Responses API sends response.completed (not response.done) as
the final event. Our wrappers and converters now emit/accept the correct
type, fixing silent usage tracking failures for native OpenAI streaming.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@SantiagoDePolonia SantiagoDePolonia self-assigned this Feb 25, 2026
@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Feb 25, 2026

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 2f357ce and e62ce42.

📒 Files selected for processing (1)
  • internal/usage/stream_wrapper.go

📝 Walkthrough

Walkthrough

Renames the streaming completion event from "response.done" to "response.completed" across providers and tests, and refactors stream usage extraction to incremental SSE parsing with buffered event processing and safer remainder bounds.

Changes

Cohort / File(s) Summary
Providers & converter updates
internal/providers/.../anthropic.go, internal/providers/responses_converter.go, internal/auditlog/stream_wrapper.go
Emit/recognize "response.completed" instead of "response.done"; update logging, SSE event formatting, and parsing to accept both new and legacy types.
Provider tests updated
internal/providers/groq/groq_test.go, internal/providers/ollama/ollama_test.go, internal/providers/openai/openai_test.go, internal/providers/xai/xai_test.go
Adjusted expected stream event name and JSON payloads from "response.done" → "response.completed" in streaming tests.
Usage stream wrapper refactor
internal/usage/stream_wrapper.go, internal/usage/constants.go
Replaced final-buffer parsing with incremental eventBuffer processing, added processCompleteEvents/parseRemainingBuffer/extractUsageFromJSON enhancements, introduced maxEventBufferRemainder constant, and removed exported SSEBufferSize.
Stream wrapper tests
internal/usage/stream_wrapper_test.go
Added comprehensive tests for incremental parsing: fragmented reads, large responses, double-close, nil/disabled logger, and usage extraction from Responses API formats.
Server/tests/e2e mocks
internal/server/handlers_test.go, tests/e2e/helpers_test.go, tests/e2e/mock_provider.go, tests/e2e/responses_test.go, tests/integration/setup_test.go
Updated mock streams and helpers to emit/recognize "response.completed" and added usage payload to final event where applicable; adjusted test expectations and comments.
Audit/log wrapper
internal/auditlog/stream_wrapper.go
Parsing updated to recognize new "response.completed" event type alongside legacy values.

Sequence Diagram(s)

(omitted)

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

  • Refactor: streaming sse #26: Changes the OpenAI→Responses stream converter; closely related to renaming the final event and adjusting converter behavior.
  • Feature: usage token tracking #43: Introduced usage extraction from Responses streaming; this PR extends/refactors that logic to support "response.completed" and incremental parsing.

Poem

🐰 I nibble bytes as events flow through,
From "done" to "completed" — a hopping review.
I stitch fragments, guard leftover crumbs,
Count tokens as each small packet comes.
Hooray — streams are snug, and logs hum too!

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Title check ⚠️ Warning The title claims the PR fixes streaming usage tracking, but the primary changes involve renaming the 'response.done' event to 'response.completed' across multiple files, with usage tracking improvements being secondary. Consider a more accurate title such as 'Rename response.done event to response.completed and improve streaming usage tracking' to reflect the actual scope of changes.
Docstring Coverage ⚠️ Warning Docstring coverage is 14.29% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (1 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
  • 📝 Generate docstrings (stacked PR)
  • 📝 Generate docstrings (commit on current branch)
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch fix/streaming-usage-tracking

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (4)
internal/providers/xai/xai_test.go (1)

649-671: 🧹 Nitpick | 🔵 Trivial

checkStream doesn't assert the response.completed event it is testing.

The mock response body at Lines 643–644 now emits response.completed, but the checkStream closure only asserts response.created, response.output_text.delta, and [DONE]. Adding a check for response.completed makes the test actually verify the renamed event survives the passthrough.

✅ Suggested addition
 if !strings.Contains(responseStr, "[DONE]") {
 	t.Error("response should end with [DONE]")
 }
+if !strings.Contains(responseStr, "response.completed") {
+	t.Error("response should contain response.completed event")
+}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@internal/providers/xai/xai_test.go` around lines 649 - 671, The test's
checkStream closure (checkStream func(t *testing.T, body io.ReadCloser))
currently asserts response.created, response.output_text.delta and [DONE] but
omits verifying the response.completed event; update the closure to also assert
that the streamed response contains "response.completed" (e.g., add an if
!strings.Contains(responseStr, "response.completed") { t.Error("response should
contain response.completed event") }) so the test confirms the renamed event
passes through.
internal/providers/groq/groq_test.go (1)

580-590: 🧹 Nitpick | 🔵 Trivial

TestStreamResponses missing assertion for response.completed.

TestStreamResponses verifies response.created and response.output_text.delta but does not assert response.completed, leaving the final event untested here. TestGroqResponsesStreamConverter covers this via the converter directly, but TestStreamResponses exercises the full StreamResponses → converter path and should assert the same terminal event for full regression coverage.

✅ Suggested addition
 if !strings.Contains(responseStr, "[DONE]") {
 	t.Error("response should end with [DONE]")
 }
+if !strings.Contains(responseStr, "response.completed") {
+	t.Error("response should contain response.completed event")
+}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@internal/providers/groq/groq_test.go` around lines 580 - 590, Add an
assertion in TestStreamResponses to check that the streaming output includes the
terminal "response.completed" event: after the existing checks on responseStr
(which currently asserts "response.created", "response.output_text.delta", and
"[DONE]"), add a strings.Contains check for "response.completed" and call
t.Error (or t.Fatalf) if it's missing so the full StreamResponses →
TestGroqResponsesStreamConverter path verifies the terminal event as well.
internal/providers/ollama/ollama_test.go (1)

640-650: 🧹 Nitpick | 🔵 Trivial

TestStreamResponses missing assertion for response.completed.

Same gap as groq/groq_test.go: TestStreamResponses verifies response.created and response.output_text.delta but omits a check for response.completed, leaving the final event untested on the full StreamResponses path.

✅ Suggested addition
 if !strings.Contains(responseStr, "[DONE]") {
 	t.Error("response should end with [DONE]")
 }
+if !strings.Contains(responseStr, "response.completed") {
+	t.Error("response should contain response.completed event")
+}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@internal/providers/ollama/ollama_test.go` around lines 640 - 650, Add an
assertion in TestStreamResponses to verify the final "response.completed" event
is present: after converting respBody to responseStr (variable responseStr) add
a check using strings.Contains(responseStr, "response.completed") and call
t.Error with a descriptive message if it's missing, mirroring the existing
checks for "response.created" and "response.output_text.delta" to ensure the
full StreamResponses path emits the completed event.
internal/providers/openai/openai_test.go (1)

624-651: ⚠️ Potential issue | 🟡 Minor

Assert response.completed explicitly in stream validation.

The fixture now uses response.completed (Line 624), but checkStream doesn’t assert it. This can let final-event regressions slip through.

Proposed assertion update
 				if !strings.Contains(responseStr, "response.output_text.delta") {
 					t.Error("response should contain response.output_text.delta event")
 				}
+				if !strings.Contains(responseStr, "response.completed") {
+					t.Error("response should contain response.completed event")
+				}
 				if !strings.Contains(responseStr, "[DONE]") {
 					t.Error("response should end with [DONE]")
 				}
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@internal/providers/openai/openai_test.go` around lines 624 - 651, The test's
stream validator (the checkStream closure in the test case) currently asserts
for "response.created", "response.output_text.delta", and "[DONE]" but doesn't
assert the new "response.completed" event; update the checkStream function used
in internal/providers/openai/openai_test.go to also verify that the response
body string contains "response.completed" (e.g., add an assertion that checks
strings.Contains(responseStr, "response.completed") and fails the test if it's
missing) alongside the existing assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@internal/usage/stream_wrapper.go`:
- Around line 68-70: The current logic uses w.eventBuffer.Reset() when
w.eventBuffer.Len() > maxEventBufferRemainder which silently drops in-flight
events (including a large response.completed). Instead of Reset(), trim the
buffer to keep only the newest maxEventBufferRemainder bytes: read b :=
w.eventBuffer.Bytes(), compute start := max(0, len(b)-maxEventBufferRemainder),
create a new bytes.Buffer (or reuse w.eventBuffer by Reset then Write) and write
b[start:], and assign that back to w.eventBuffer; apply this change at the spots
referencing w.eventBuffer and maxEventBufferRemainder (the branches currently
calling w.eventBuffer.Reset()) so partial events are preserved rather than
discarded.

---

Outside diff comments:
In `@internal/providers/groq/groq_test.go`:
- Around line 580-590: Add an assertion in TestStreamResponses to check that the
streaming output includes the terminal "response.completed" event: after the
existing checks on responseStr (which currently asserts "response.created",
"response.output_text.delta", and "[DONE]"), add a strings.Contains check for
"response.completed" and call t.Error (or t.Fatalf) if it's missing so the full
StreamResponses → TestGroqResponsesStreamConverter path verifies the terminal
event as well.

In `@internal/providers/ollama/ollama_test.go`:
- Around line 640-650: Add an assertion in TestStreamResponses to verify the
final "response.completed" event is present: after converting respBody to
responseStr (variable responseStr) add a check using
strings.Contains(responseStr, "response.completed") and call t.Error with a
descriptive message if it's missing, mirroring the existing checks for
"response.created" and "response.output_text.delta" to ensure the full
StreamResponses path emits the completed event.

In `@internal/providers/openai/openai_test.go`:
- Around line 624-651: The test's stream validator (the checkStream closure in
the test case) currently asserts for "response.created",
"response.output_text.delta", and "[DONE]" but doesn't assert the new
"response.completed" event; update the checkStream function used in
internal/providers/openai/openai_test.go to also verify that the response body
string contains "response.completed" (e.g., add an assertion that checks
strings.Contains(responseStr, "response.completed") and fails the test if it's
missing) alongside the existing assertions.

In `@internal/providers/xai/xai_test.go`:
- Around line 649-671: The test's checkStream closure (checkStream func(t
*testing.T, body io.ReadCloser)) currently asserts response.created,
response.output_text.delta and [DONE] but omits verifying the response.completed
event; update the closure to also assert that the streamed response contains
"response.completed" (e.g., add an if !strings.Contains(responseStr,
"response.completed") { t.Error("response should contain response.completed
event") }) so the test confirms the renamed event passes through.

ℹ️ Review info

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between d5015be and 2f357ce.

📒 Files selected for processing (15)
  • internal/auditlog/stream_wrapper.go
  • internal/providers/anthropic/anthropic.go
  • internal/providers/groq/groq_test.go
  • internal/providers/ollama/ollama_test.go
  • internal/providers/openai/openai_test.go
  • internal/providers/responses_converter.go
  • internal/providers/xai/xai_test.go
  • internal/server/handlers_test.go
  • internal/usage/constants.go
  • internal/usage/stream_wrapper.go
  • internal/usage/stream_wrapper_test.go
  • tests/e2e/helpers_test.go
  • tests/e2e/mock_provider.go
  • tests/e2e/responses_test.go
  • tests/integration/setup_test.go
💤 Files with no reviewable changes (1)
  • internal/usage/constants.go

When the SSE event buffer exceeds maxEventBufferRemainder, keep the
newest bytes rather than resetting entirely. This preserves the tail of
in-flight partial events (where usage data in JSON is most likely found)
instead of silently dropping them.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@SantiagoDePolonia SantiagoDePolonia merged commit e034f66 into main Feb 25, 2026
11 checks passed
@SantiagoDePolonia SantiagoDePolonia deleted the fix/streaming-usage-tracking branch March 22, 2026 14:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant