Skip to content

byok: keep streamed token usage when a provider omits the final line break - #330547

Open
Vritant Bhardwaj (vritant24) wants to merge 1 commit into
mainfrom
agents/investigate-issue-329436-root-cause
Open

byok: keep streamed token usage when a provider omits the final line break#330547
Vritant Bhardwaj (vritant24) wants to merge 1 commit into
mainfrom
agents/investigate-issue-329436-root-cause

Conversation

@vritant24

Copy link
Copy Markdown
Member

Summary

  • Restores token usage for custom OpenAI-compatible (BYOK) models in streaming mode, so the context window indicator stays visible instead of disappearing.
  • Fixes usage being logged as undefined when a provider ends its stream immediately after the final usage event, without a terminating line break or [DONE].
  • Scoped to endpoints that talk directly to a provider. Requests that go through the Copilot proxy keep their existing stream handling.

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 its data: prefix, so parsing always failed and the usage payload was discarded.

The completion itself was unaffected, because content and finish_reason arrive in earlier, well-terminated events. Only the trailing usage event sat in the vulnerable position. That matches the report exactly: text streams normally, usage is undefined, 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

  • Proxy-backed requests resolve the factory to the pre-existing behavior, so their stream handling is unchanged.
  • Streams that already end with [DONE] finish before the appended newline is ever read, so well-formed responses of either kind behave identically.
  • The appended newline is emitted only on normal stream completion, not on cancellation or abort, so cancellation semantics are preserved.
  • The transform participates in the existing stream wrapper, so cancellation still propagates to the underlying network stream.
  • A truncated or malformed trailing event still fails to parse and is reported through existing diagnostics; the completion is still delivered rather than dropped.
  • Responses API, Anthropic Messages API, and non-streaming requests are untouched and continue through their existing processors.

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.

…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 AI balanced review requested due to automatic review settings August 12, 2026 21:04

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

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.

Custom OpenAI-compatible model streaming response usage metadata shows as "undefined" in Copilot chat log

3 participants