Fix bedrock mantle streaming 5871

[zbloss]

Summary

bedrock_mantle streaming chat completions never finish: every chunk is delivered, including the terminal chunk with finish_reason and usage, and then the response stalls - data: [DONE] is never sent and the HTTP body is never terminated, leaking one goroutine and one upstream socket per request. Two defects combine to cause this, and this PR fixes both.

Changes

• core/providers/utils/utils.go - ReleaseStreamingResponse drained the upstream body with an unbounded io.Copy. BuildStreamingClient sets ReadTimeout = 0, and the drain runs after the SSE loop's own idle-timeout reader has already been stopped, so nothing bounds it - an upstream that holds its chunked body open past the terminal event (as bedrock_mantle does) blocks the copy forever, parking the provider goroutine in this defer before the defer that closes the response channel runs. The drain now runs in its own goroutine and races it against GetStreamIdleTimeout(ctx); on timeout the connection is force-closed with ErrStreamIdleTimeout (reusing the same sentinel NewIdleTimeoutReader's timer already uses for this failure mode, rather than adding a duplicate) so fasthttp drops it instead of returning it to the idle pool.
• core/providers/utils/utils.go - ProviderSendsDoneMarker didn't list schemas.BedrockMantle, so the SSE loop waited indefinitely for a [DONE] marker Mantle never sends. Added it alongside the existing schemas.Bedrock case.
• Design decision: on the forced-close timeout path, resp.CloseBodyStream()/fasthttp.ReleaseResponse(resp) are deliberately not called — the drain goroutine may still be mid-Read on the body stream, and a second concurrent close/release would double-Put fasthttp's pooled reader/requestStream, corrupting a later, unrelated request. resp is intentionally leaked to GC in that rare path, the same trade-off already used elsewhere in this function when close ownership is uncertain.

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 version
cd core
go test ./providers/utils/... -race -v -run 'TestReleaseStreamingResponse|TestProviderSendsDoneMarker'
go test ./...

Expected: 1TestReleaseStreamingResponse_BoundsDrainAndDropsConnection1 completes in well under the 2s bound instead of hanging (reproduces the reported hang against a body stream that never terminates), TestReleaseStreamingResponse_NormalDrainReleasesWithoutForcingClose and TestProviderSendsDoneMarker (including the new bedrock_mantle case) pass, and the full go test ./... run shows no new failures (providers/bedrock and internal/mcptests have pre-existing, unrelated failures on dev).

Manual repro against a mock upstream (from the linked issue): run a TLS server aliased to bedrock-mantle.us-east-1.api.aws that sends a few delta.content frames, a terminal frame with finish_reason, data: [DONE], and then never sends the closing 0\r\n\r\n chunk. Configure a bedrock_mantle key against it and stream a request through /v1/chat/completions. Before this change the client hangs indefinitely; after, data: [DONE] and the stream end arrive within the configured stream_idle_timeout_in_seconds.

No new configs or environment variables were added.

Breaking changes

• Yes
• No

If yes, describe impact and migration instructions.

Related issues

Closes #5871
Relates to #2910 (same missing-[DONE]-marker defect, reported for custom providers; that report's hang self-heals at the stream idle timeout, this one didn't because of the unbounded drain)

Security considerations

None. The change only affects timing/cleanup of an already-authenticated provider connection: it bounds how long a stalled upstream can hold a goroutine and socket open, and ensures a connection that failed to terminate cleanly is dropped rather than returned to the pool for reuse (avoiding stale-connection reuse, not introducing it). No auth, secrets, or PII handling is touched.

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

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 4800aae8-2427-433d-9ee2-d8600f63d2f1

📥 Commits

Reviewing files that changed from the base of the PR and between 259ca31 and b6b04ce.

📒 Files selected for processing (1)

• core/providers/utils/idle_timeout_reader_test.go

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved streaming response cleanup to prevent connections from hanging indefinitely.
  • Streaming responses that stall now time out and close automatically.
  • Updated stream completion handling for Bedrock Mantle.
  • Preserved normal cleanup behavior for responses that finish successfully.

Walkthrough

The PR adds Bedrock Mantle completion-marker handling and bounds streaming-response draining by the configured idle timeout. It adds tests for provider marker behavior, stalled stream closure, timeout enforcement, and normal response cleanup.

Changes

Streaming utility behavior

Layer / File(s)	Summary
Provider completion marker handling core/providers/utils/utils.go, core/providers/utils/utils_test.go	ProviderSendsDoneMarker now excludes Bedrock Mantle. Tests cover marker expectations for supported providers.
Bounded streaming response cleanup core/providers/utils/utils.go, core/providers/utils/idle_timeout_reader_test.go	ReleaseStreamingResponse drains asynchronously and force-closes stalled streams with ErrStreamIdleTimeout. Tests cover timeout behavior, close errors, and normal drain-and-release cleanup.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ReleaseStreamingResponse
  participant GetStreamIdleTimeout
  participant ResponseBody
  ReleaseStreamingResponse->>GetStreamIdleTimeout: Read configured idle timeout
  ReleaseStreamingResponse->>ResponseBody: Start asynchronous drain
  GetStreamIdleTimeout-->>ReleaseStreamingResponse: Return timeout
  ReleaseStreamingResponse->>ResponseBody: CloseWithError(ErrStreamIdleTimeout) when drain stalls
  ResponseBody-->>ReleaseStreamingResponse: Unblock read and report close

Loading

Possibly related issues

• [Bug]: AWS Bedrock Mantle streaming is broken #5871 — The PR updates both Bedrock Mantle marker handling and bounded streaming-response draining described by the issue.

Possibly related PRs

• maximhq/bifrost#5568 — Both changes modify streaming [DONE] marker behavior.

Suggested reviewers: tejasghatte, akshaydeo

🚥 Pre-merge checks | ✅ 4 | ❌ 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 (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title identifies the Bedrock Mantle streaming fix, which matches the primary hang addressed by the pull request.
Description check	✅ Passed	The description explains the problem, changes, tests, impact, related issues, and security considerations with appropriate checklist selections.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"

---

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

• X
• Mastodon
• Reddit
• LinkedIn

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

[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/idle_timeout_reader_test.go`:
- Around line 758-783: Strengthen the test around ReleaseStreamingResponse by
recording when body.CloseWithError is invoked and asserting that it occurs
within a CI-safe upper bound derived from the configured 30ms stream idle
timeout, rather than only checking it completes before 2 seconds. Preserve the
existing ErrStreamIdleTimeout assertion and use the blocking body’s signaling to
measure the forced-close timing deterministically.
- Around line 801-811: Update countingStreamCloser and the normal-path test
around ReleaseStreamingResponse to record the error argument passed to
CloseWithError, then assert that the recorded error is nil after cleanup. Keep
the existing close-count assertion and ensure the test distinguishes the normal
nil-error close from the idle-timeout forced-close branch.

🪄 Autofix

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: CHILL

Plan: Pro Plus

Run ID: 1494350d-a489-4370-8014-7b6d65c540cc

📥 Commits

Reviewing files that changed from the base of the PR and between e575bab and 259ca31.

📒 Files selected for processing (3)

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

[zbloss]

@akshaydeo @TejasGhatte do you have an estimate on when you will review these changes? I'm wondering if I need to fork this branch and patch my local deployments.

[TejasGhatte]

@akshaydeo @TejasGhatte do you have an estimate on when you will review these changes? I'm wondering if I need to fork this branch and patch my local deployments.

Hey @zbloss approved it. @akshaydeo will merge soon

The merge-base changed after approval.

[zbloss]

@akshaydeo Do I need to fix these merge conflicts?

[zbloss]

Looks like this was addressed in 1.7.11? https://github.com/maximhq/bifrost/releases#release-core/v1.7.11

[zbloss]

@TejasGhatte @akshaydeo I think we still need either this PR or another that changes

bifrost/core/providers/utils/utils.go

Lines 3183 to 3185 in 5bb5c65

	case schemas.Cerebras, schemas.Perplexity, schemas.HuggingFace, schemas.Bedrock:
	// Cerebras, Perplexity, HuggingFace, and Bedrock mantle don't send [DONE] marker, ends stream after finish_reason
	return false

Bedrock Mantle is not listed here even though it's in the code comments, which is causing our bedrock mantle streams to hang open waiting on a [DONE] marker that never arrives

I think this should be

case schemas.Cerebras, schemas.Perplexity, schemas.HuggingFace, schemas.Bedrock, schemas.BedrockMantle:
		// Cerebras, Perplexity, HuggingFace, and Bedrock mantle don't send [DONE] marker, ends stream after finish_reason
		return false

[zbloss]

This looks related: #6021

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
7 out of 9 committers have signed the CLA.

✅ jeremym-tanium
✅ CMWR421
✅ Madhuvod
✅ AdityaPainuli
✅ TransactCharlie
✅ zbloss
✅ akshaydeo
❌ roroghost17
❌ Pratham-Mishra04
_{You have signed the CLA already but the status is still pending? Let us recheck it.}