fix: set ConnectionClosed flag before closing stream on cancellation

[TejasGhatte]

Summary

Fixes a race condition in SetupStreamCancellation where BifrostContextKeyConnectionClosed was set after calling Close() on the body stream. Because Close() immediately unblocks any in-progress Read (which may panic due to a force-closed connection), the recover block in idleTimeoutReader.Read could run before the flag was set — causing it to re-panic instead of returning ErrStreamClosed.

Additionally, idleTimeoutReader.Read was returning (0, nil) when the connection was already marked closed, which is incorrect. It now returns the appropriate closed-stream error via closedReadError().

Changes

• SetupStreamCancellation now sets BifrostContextKeyConnectionClosed before calling Close() or CloseWithError() in all four call sites, eliminating the race window where a panicking Read could recover before the flag was visible.
• idleTimeoutReader.Read now returns r.closedReadError() instead of (0, nil) when the connection is already closed, ensuring callers receive a meaningful error rather than a silent empty read.
• Added syncedPanicBody, a deterministic test helper that reproduces the exact race: Close() triggers a panic in Read and then blocks, holding SetupStreamCancellation inside Close() so the flag is guaranteed to be unset when the recover block runs under the unfixed code.
• Added TestSetupStreamCancellation_NoPanicOnCancelledContext which fails against the unfixed code and passes after the fix.

Type of change

• Bug fix

Affected areas

• Core (Go)
• Transports (HTTP)

How to test

go test ./core/providers/utils/... -race -count=1 -run TestSetupStreamCancellation_NoPanicOnCancelledContext
go test ./core/providers/utils/... -race -count=1

The new test should pass without any detected data races or re-panics.

Breaking changes

• No

Security considerations

None. This is a stability fix for stream teardown on context cancellation.

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]

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Resolved potential panic when cancelling in-flight stream operations during context cancellation
  • Improved error handling for reads on closed streams to properly return error responses

• Tests

  • Added regression test for context cancellation behavior with active stream operations

Walkthrough

This PR prevents panic during context cancellation in stream handling by establishing flag-before-close ordering: SetupStreamCancellation now sets BifrostContextKeyConnectionClosed before closing the upstream stream in both cancellation paths, and idleTimeoutReader.Read returns a proper closed error instead of nil. A deterministic regression test validates the fix.

Changes

Stream Cancellation Safety

Layer / File(s)	Summary
Flag-before-close in stream cancellation and error handling core/providers/utils/utils.go	SetupStreamCancellation sets BifrostContextKeyConnectionClosed to true before closing the stream in both the ctx.Done() branch and the case <-done branch when context is cancelled. idleTimeoutReader.Read now returns closedReadError() instead of nil when the stream is already marked closed, ensuring callers observe the proper closed state.
Regression test with synchronization helper core/providers/utils/idle_timeout_reader_test.go	syncedPanicBody provides a deterministic io.ReadCloser that blocks in Read until Close() is called, then panics. TestSetupStreamCancellation_NoPanicOnCancelledContext cancels the context during in-flight reads and verifies that reads do not panic and return ErrStreamClosed instead.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• maximhq/bifrost#3677: Both PRs modify core/providers/utils/utils.go around the idleTimeoutReader closed-read and no-panic behavior with test coverage.
• maximhq/bifrost#3582: Both PRs adjust stream-cancellation behavior by checking/setting schemas.BifrostContextKeyConnectionClosed before close handling to prevent bad cleanup and panics.
• maximhq/bifrost#3495: Both PRs modify SetupStreamCancellation and idleTimeoutReader close behavior in the stream-timeout/cancellation path.

Suggested reviewers

• danpiths
• akshaydeo

Poem

🐰 When contexts cancel in the streaming night,
A flag set first prevents the panic fright,
No double-close, no re-entrant despair—
Just tidy errors floating through the air!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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 clearly and concisely describes the main fix: setting the ConnectionClosed flag before closing the stream during cancellation.
Description check	✅ Passed	The PR description is comprehensive and well-structured, covering summary, changes, type, affected areas, testing instructions, and security considerations.
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 docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 05-25-fix_set_connectionclosed_flag_before_closing_stream_on_cancellation

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"

---

_{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.}

[TejasGhatte]

• fix: set ConnectionClosed flag before closing stream on cancellation #3733 👈 (View in Graphite)
• dev

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

[greptile-apps]

Confidence Score: 5/5

Safe to merge — the changes are narrow, well-tested, and address a real stream teardown crash with no API surface changes.

Both changes are targeted: reordering two lines to fix a flag-visibility race, and replacing a silent (0, nil) return with a proper error. The new test deterministically reproduces the race condition and the defer ordering in the test is correct. No existing callers are broken since ErrStreamClosed/ErrStreamIdleTimeout were already the established error values for this path.

No files require special attention.

Important Files Changed

Filename	Overview
core/providers/utils/utils.go	Fixes the race in SetupStreamCancellation (flag now set before Close/CloseWithError in all 4 sites) and fixes idleTimeoutReader.Read returning (0, nil) when connection is already closed — both are targeted, correct changes.
core/providers/utils/idle_timeout_reader_test.go	Adds syncedPanicBody (deterministic race reproducer) and TestSetupStreamCancellation_NoPanicOnCancelledContext; defer ordering and channel synchronization are correct and the test will reliably fail against the old code.

_{Reviews (1): Last reviewed commit: "fix: set ConnectionClosed flag before cl..." | Re-trigger Greptile}

[coderabbitai]

🧹 Nitpick comments (1)

core/providers/utils/idle_timeout_reader_test.go (1)

426-488: ⚡ Quick win

Add a companion regression for the done + cancelled-context path.

This test only drives the <-ctx.Done() branch. The PR also changes the <-done branch in SetupStreamCancellation, and that ordering bug is subtle enough that it’s worth pinning with a second test as well.

🤖 Prompt for 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.

In `@core/providers/utils/idle_timeout_reader_test.go` around lines 426 - 488, Add
a second regression test mirroring
TestSetupStreamCancellation_NoPanicOnCancelledContext but exercising the
"<-done" branch: create a test (e.g.,
TestSetupStreamCancellation_NoPanicOnDoneClose) that uses NewIdleTimeoutReader,
SetupStreamCancellation and the same synced panic body, spawn a Read() goroutine
that recovers and reports panics/errors, then trigger the done path by closing
the body's done channel (instead of cancelling the context) and assert that Read
does not re-panic and returns ErrStreamClosed; reference the same helper symbols
(NewIdleTimeoutReader, SetupStreamCancellation, the synced panic body fields
like allowReturn/done) so the subtle ordering bug for the done-branch is pinned.

🤖 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.

Nitpick comments:
In `@core/providers/utils/idle_timeout_reader_test.go`:
- Around line 426-488: Add a second regression test mirroring
TestSetupStreamCancellation_NoPanicOnCancelledContext but exercising the
"<-done" branch: create a test (e.g.,
TestSetupStreamCancellation_NoPanicOnDoneClose) that uses NewIdleTimeoutReader,
SetupStreamCancellation and the same synced panic body, spawn a Read() goroutine
that recovers and reports panics/errors, then trigger the done path by closing
the body's done channel (instead of cancelling the context) and assert that Read
does not re-panic and returns ErrStreamClosed; reference the same helper symbols
(NewIdleTimeoutReader, SetupStreamCancellation, the synced panic body fields
like allowReturn/done) so the subtle ordering bug for the done-branch is pinned.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: caf8c37c-252c-4a13-8f8d-6e8638b04738

📥 Commits

Reviewing files that changed from the base of the PR and between 584816b and 282430b.

📒 Files selected for processing (2)

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

[akshaydeo]

Merge activity

• May 25, 4:28 PM UTC: A user started a stack merge that includes this pull request via Graphite.
• May 25, 4:29 PM UTC: @akshaydeo merged this pull request with Graphite.