byok: keep streamed token usage when a provider omits the final line break - #330547
Open
Vritant Bhardwaj (vritant24) wants to merge 1 commit into
Open
byok: keep streamed token usage when a provider omits the final line break#330547Vritant Bhardwaj (vritant24) wants to merge 1 commit into
Vritant Bhardwaj (vritant24) wants to merge 1 commit into
Conversation
…break Custom OpenAI-compatible endpoints request token usage through `stream_options.include_usage`, so the server sends usage in a final SSE event. Some of these servers close the connection right after that event without terminating the line or sending `[DONE]`. The parser kept the unterminated line buffered and dropped it, so usage was reported as undefined and the context window indicator disappeared. Direct-to-provider endpoints now use an SSE processor that tolerates a missing final line break. The CAPI path keeps its existing processor. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot started reviewing on behalf of
Vritant Bhardwaj (vritant24)
August 12, 2026 21:06
View session
Contributor
There was a problem hiding this comment.
Pull request overview
Preserves streamed token usage for direct OpenAI-compatible BYOK endpoints when the final SSE event lacks a terminating newline.
Changes:
- Adds an opt-in SSE processor that flushes the final buffered line.
- Enables tolerant processing for direct BYOK endpoints while preserving proxy behavior.
- Adds processor-level and endpoint-level regression coverage.
Show a summary per file
| File | Description |
|---|---|
extensions/copilot/src/platform/networking/node/stream.ts |
Adds trailing-line flushing. |
extensions/copilot/src/platform/endpoint/node/chatEndpoint.ts |
Introduces an overridable processor factory. |
extensions/copilot/src/extension/byok/node/openAIEndpoint.ts |
Enables tolerant processing for BYOK. |
extensions/copilot/src/platform/endpoint/test/node/stream.sseProcessor.spec.ts |
Covers framing and malformed-input behavior. |
extensions/copilot/src/extension/byok/node/test/openAIEndpoint.spec.ts |
Verifies end-to-end usage retention. |
Review details
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
- Files reviewed: 5/5 changed files
- Comments generated: 0
- Review effort level: Balanced
Paul (pwang347)
approved these changes
Aug 12, 2026
Vritant Bhardwaj (vritant24)
enabled auto-merge
August 12, 2026 23:24
Vritant Bhardwaj (vritant24)
disabled auto-merge
August 12, 2026 23:24
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
undefinedwhen a provider ends its stream immediately after the final usage event, without a terminating line break or[DONE].Fixes #329436
Technical context for AI-assisted review
Intent and previous behavior
Client-side BYOK endpoints ask for token accounting by setting
stream_options.include_usage, which makes the server emit a final, usage-only SSE event after the content events. The SSE reader splits incoming network chunks on newlines and holds any unterminated remainder in a buffer for the next chunk. When a provider closed the connection directly after that final event without terminating the line and without a[DONE]sentinel, the event stayed in that buffer. End-of-stream handling then tried to parse the buffer as bare JSON while it still carried itsdata:prefix, so parsing always failed and the usage payload was discarded.The completion itself was unaffected, because content and
finish_reasonarrive in earlier, well-terminated events. Only the trailing usage event sat in the vulnerable position. That matches the report exactly: text streams normally, usage isundefined, and turning streaming off works because the non-streaming path parses one complete JSON body and never performs line framing. Because this usage value is the sole source for the BYOK context window indicator, losing it removes the indicator entirely.Implementation
The stream reader gains a second construction path that appends one newline after the source stream ends, letting the existing parser surface the buffered final event through its normal per-event handling. No parsing, buffering, or event-dispatch logic changed.
Endpoint wiring selects between the two construction paths through an overridable factory on the chat endpoint. The base endpoint returns the existing factory, and the direct-to-provider endpoint returns the tolerant one. The endpoint's branch that decides which processor applies still owns that decision and calls the same response processor as before, now passing the factory. Keeping the decision in one place avoids restating the API-selection conditions in the subclass, where they could drift if another API path is added later.
Behavior and constraints
[DONE]finish before the appended newline is ever read, so well-formed responses of either kind behave identically.Review context
The reported symptom is reproduced by a regression test, and existing coverage confirms usage is retained for well-framed streams. The provider-side trigger is inferred rather than captured from the reporter's traffic, since the report does not include raw response bytes. The inference follows from the fact that well-framed streams already report usage correctly, so the failing stream must be framed differently; an unterminated final event is the realistic form of that difference for minimal server implementations.
One tradeoff worth confirming: the factory seam is exposed as a property returning a static function reference. Both factories are pure and ignore
this, so this is safe today, but a future factory that relied on instance state would need binding.