fix: drain non sse stream reader

[TejasGhatte]

Summary

Some OpenAI-compatible backends return valid SSE frames without a Content-Type: text/event-stream header. The previous DrainNonSSEStreamResponse helper would unconditionally drain and discard the response body in this case, causing the downstream SSE parser to receive an empty stream. This PR introduces a reader-preserving variant that peeks at the first bytes of the stream to detect SSE field prefixes (data:, event:, id:, retry:, :, or leading newlines) before deciding whether to drain or pass the reader through intact.

Changes

• Introduced DrainNonSSEStreamReader(resp, reader) which accepts an io.Reader (e.g. a decompressed stream) and returns a potentially buffered reader alongside a drained boolean, preserving the stream when it looks like SSE even if the content type header is absent.
• DrainNonSSEStreamResponse is retained as a thin wrapper delegating to DrainNonSSEStreamReader for backward compatibility.
• All streaming handlers in the OpenAI provider (text completion, chat completion, responses, speech, transcription, image generation, image edit) now use DrainNonSSEStreamReader and reassign the reader from its return value so the buffered peek bytes are not lost.
• Added looksLikeSSEPrefix to perform a case-sensitive, 16-byte peek-based heuristic for SSE field prefixes.
• Added tests covering: SSE without content type remains readable, gzip-compressed SSE without content type remains readable, JSON without content type is drained, and uppercase SSE-like prefixes are treated as non-SSE.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go test ./core/providers/utils/... -run TestDrainNonSSEStreamReader
go test ./...

Expected: all new TestDrainNonSSEStreamReader_* tests pass, and existing streaming tests remain green. To validate end-to-end, route a streaming request through a backend that returns SSE without Content-Type: text/event-stream and confirm the response is streamed correctly rather than returning an error.

Breaking changes

• Yes
• No

Related issues

Security considerations

None. The peek reads at most 16 bytes from the stream and does not log or expose any content.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

Summary by CodeRabbit

• Bug Fixes

  • Improved streaming across text, chat, responses, speech, transcription, image generation and edit flows to detect SSE-like streams, avoid prematurely draining them, and prevent stream hangs or unexpected termination.

• Tests

  • Added and updated unit tests for SSE detection, compressed-stream handling, fragmented/tiny-prefix delivery, and correct draining behavior for non-SSE payloads.

[TejasGhatte]

• fix: drain non sse stream reader #3956 👈 (View in Graphite)
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Introduce DrainNonSSEStreamReader to peek and preserve SSE-like streams when Content-Type is missing/incorrect, replace prior response-only draining with reader-aware draining across seven OpenAI streaming goroutines, and add unit tests covering plain, gzip, short-read SSE detection and non-SSE draining.

Changes

SSE Stream Detection Refactoring

Layer / File(s)	Summary
Buffered SSE detection utility core/providers/utils/utils.go	Add DrainNonSSEStreamReader(resp, reader) and SSE-prefix helpers to peek initial bytes and decide whether to preserve the reader or drain to io.Discard. Includes bufio import.
Unit tests for detection core/providers/utils/utils_test.go	Add tests verifying SSE-like payloads (plain, gzip, short-read, live-stream) without Content-Type are preserved, JSON-like payloads are drained, uppercase DATA: is drained, and include shortReadReader and timing tests.
OpenAI streaming endpoint integration core/providers/openai/openai.go	Update text completion, chat completion, responses, speech, transcription, image generation, and image edit streaming goroutines to call DrainNonSSEStreamReader(resp, reader) and branch on the returned drained flag instead of DrainNonSSEStreamResponse.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested reviewers

• danpiths
• akshaydeo

Poem

A rabbit peeks where streams begin,
I nibble prefixes, keep what's thin.
When "data:" whispers on the line,
I hold the thread and let it shine.
Hooray for streaming—🐇🌿

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'fix: drain non sse stream reader' clearly and concisely describes the main bug fix, matching the core change of updating non-SSE stream handling logic.
Description check	✅ Passed	The PR description covers all required template sections: summary explaining the problem, detailed changes, type of change marked, affected areas selected, testing instructions provided, breaking changes addressed, and checklist items marked appropriately.
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.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 06-01-fix_drain_non_sse_stream_reader

---

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

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

tejas ghatte seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/providers/utils/utils.go`:
- Around line 1003-1005: DrainNonSSEStreamResponse currently discards the reader
returned by DrainNonSSEStreamReader, which lets Peek(16) buffer bytes in a
throwaway bufio.Reader and causes subsequent reads from resp.BodyStream() to
miss those bytes; instead call DrainNonSSEStreamReader(resp, resp.BodyStream()),
capture the returned io.Reader (the possibly wrapped/buffered reader) and
restore it onto the response (e.g. via resp.SetBodyStream(returnedReader, -1) or
the equivalent in fasthttp) before returning the drained bool so callers keep
the full stream.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 4bbb409a-ec9b-4b1f-9101-3dc6197562aa

📥 Commits

Reviewing files that changed from the base of the PR and between d4c96b8 and 435e381.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/utils/utils.go
• core/providers/utils/utils_test.go

[greptile-apps]

Confidence Score: 5/5

Safe to merge; the core logic is well-tested, all seven streaming handlers have been updated consistently, and no pre-existing callers of the removed function remain.

The change is narrowly scoped: a new peek-based heuristic replaces an unconditional drain, and the returned reader is correctly reassigned in every handler. The partial-prefix match behaviour is intentional and pinned by dedicated tests. No regressions to non-SSE paths were found, and the idle-timeout / decompression / cancellation chain interacts with the new bufio wrapper correctly on both the SSE and drain code paths.

core/providers/utils/utils.go — the peekHasPrefix partial-match design is intentional but undocumented in the code itself; worth a short inline comment to prevent future misreads.

Important Files Changed

Filename	Overview
core/providers/utils/utils.go	Replaces DrainNonSSEStreamResponse with DrainNonSSEStreamReader; introduces hasSSEPrefix and peekHasPrefix for a non-blocking 1-byte peek heuristic. Logic is intentional but partial-prefix matching (when fewer bytes than the field name are buffered) can produce false positives for non-SSE bodies starting with d/e/i/r.
core/providers/openai/openai.go	All seven streaming handlers now use DrainNonSSEStreamReader and reassign the returned reader, ensuring the peek bytes buffered during SSE detection are preserved for the downstream SSE scanner. Mechanically straightforward; no logic regressions visible.
core/providers/utils/utils_test.go	Replaces two simple SSE/non-SSE drain tests with a comprehensive suite: SSE without content-type, gzip-compressed SSE, JSON drain, uppercase rejection, short-read delivery, fragmented single-byte prefix, and an end-to-end Codex reproduction test. Coverage is thorough.

_{Reviews (6): Last reviewed commit: "fix: drain non sse stream reader" | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/providers/utils/utils.go`:
- Around line 1008-1016: The probe in DrainNonSSEStreamReader currently does a
single reader.Read into a 16-byte slice which can misclassify SSE when readers
return short reads; change the logic to wrap the incoming io.Reader in a
bufio.Reader, use bufio.Reader.Peek(16) (handling smaller available bytes
without consuming) to check looksLikeSSEPrefix, and if it looks like SSE return
the bufio.Reader and false, otherwise drain the bufio.Reader (io.Copy to
io.Discard) and return nil,true; update references to the original reader
variable to use the buffered reader and keep uses of looksLikeSSEPrefix
unchanged. Also add a unit test that supplies an io.Reader implementation that
deliberately returns short reads (n < len and err == nil) to ensure the SSE
prefix split across reads is detected correctly.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 268f1b8a-2d34-4d5f-9708-690d07106f35

📥 Commits

Reviewing files that changed from the base of the PR and between 435e381 and 909f828.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/utils/utils.go
• core/providers/utils/utils_test.go

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/providers/utils/utils.go`:
- Around line 1009-1015: The current DrainNonSSEStreamReader uses
bufio.Reader.Peek(16) which can block; change DrainNonSSEStreamReader to probe
incrementally (e.g., Peek(1) first), and only call Peek(n) for longer tokens
when the first byte matches SSE candidate prefixes (':' or 'i' etc.), using
looksLikeSSEPrefix or similar logic to decide when to peek further and avoid a
fixed Peek(16); after detecting non-SSE, drain the rest without blocking. Also
add a regression test that feeds the reader small fragments like ":\n\n" or
"id:" with a pause before 16 bytes are available and assert
DrainNonSSEStreamReader returns promptly and the underlying stream remains
readable.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 3378b7dd-b792-4288-b66c-0d7f8262f7f5

📥 Commits

Reviewing files that changed from the base of the PR and between 909f828 and 89c2a42.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/utils/utils.go
• core/providers/utils/utils_test.go

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/providers/utils/utils.go`:
- Around line 997-1016: Reintroduce the exported compatibility wrapper
DrainNonSSEStreamResponse that preserves the original public API and simply
forwards to the new helper DrainNonSSEStreamReader; implement
DrainNonSSEStreamResponse with the original signature, call
DrainNonSSEStreamReader internally, and translate/return values exactly as
before (ensuring nil/reader behavior is preserved), so callers of
DrainNonSSEStreamResponse keep working; keep existing helpers like hasSSEPrefix
and sseInitialBufSize unchanged.
- Around line 1018-1041: The peekHasPrefix implementation can block because
bufio.Reader.Peek(len(prefix)) waits for all bytes; update peekHasPrefix to use
reader.Buffered() to only check min(n, len(prefix)) bytes and compare peeked
bytes to prefix[:n] so fragmented SSE names like "d" vs "data:" don't stall
hasSSEPrefix/DrainNonSSEStreamReader; adjust hasSSEPrefix to rely on this
non-blocking probe for cases 'd','e','i','r' and add a regression test that uses
io.Pipe to write only the first byte of "data:" and "event:" then assert
DrainNonSSEStreamReader returns promptly without draining the rest.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 5763115f-3ce7-4761-8171-81d72669099c

📥 Commits

Reviewing files that changed from the base of the PR and between 89c2a42 and 1f19cc0.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/utils/utils.go
• core/providers/utils/utils_test.go

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@core/providers/utils/utils_test.go`:
- Around line 824-893: The 200ms deadline in
TestDrainNonSSEStreamReader_TinyOpenSSEPrefixReturnsPromptly is too tight and
can flake on CI; update the test to use a more forgiving timeout (e.g., replace
time.After(200 * time.Millisecond) with a larger constant such as time.After(2 *
time.Second)) or, preferably, convert the timing assertion to a handshake-based
synchronization using the existing writeErr or a dedicated done channel so the
goroutine signals readiness deterministically; apply the same change to the
other promptness test referencing DrainNonSSEStreamReader to avoid flaky
failures.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 44769f20-592a-4d8d-bdc6-9f00ed27623d

📥 Commits

Reviewing files that changed from the base of the PR and between 1f19cc0 and fe25cc8.

📒 Files selected for processing (3)

• core/providers/openai/openai.go
• core/providers/utils/utils.go
• core/providers/utils/utils_test.go

[akshaydeo]

Merge activity

• Jun 2, 7:03 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Jun 2, 7:03 AM UTC: @akshaydeo merged this pull request with Graphite.