fix: prevent send on closed channel panic in provider queue shutdown --skip-pipeline

[Pratham-Mishra04]

Summary

Fixes a race condition in provider queue shutdown that caused "send on closed channel" panics in production. The issue occurred when producers passed the isClosing() check but then attempted to send to a queue that was closed before they reached the select statement.

Changes

• Removed queue channel closure: Queue channels are never closed to prevent "send on closed channel" panics
• Updated worker exit mechanism: Workers now exit via the done channel signal instead of waiting for queue closure
• Enhanced shutdown handling: Workers drain remaining buffered requests and send shutdown errors when done is signaled
• Added producer re-routing: Stale producers can transparently re-route to new queues during UpdateProvider
• Improved error handling: Added rollback logic for failed provider updates with proper cleanup
• Enhanced transfer logic: Buffered requests are transferred before signaling shutdown to ensure they reach new workers
• Added comprehensive tests: Race condition demonstration and validation of the fix across multiple scenarios

Type of change

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

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (Next.js)
• Docs

How to test

Run the new race condition test to verify the fix:

go test -run TestProviderQueue_SendOnClosedChannel_Race ./core -v

Run the comprehensive provider lifecycle tests:

go test -run TestProviderQueue ./core -v
go test -run TestUpdateProvider ./core -v
go test -run TestRemoveProvider ./core -v

Run the full test suite to ensure no regressions:

go test ./...

Screenshots/Recordings

N/A

Breaking changes

• Yes
• No

Related issues

Fixes production panics related to concurrent provider queue operations during shutdown/updates.

Security considerations

None - this is an internal concurrency fix that doesn't affect external interfaces or data handling.

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

[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.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Pratham-Mishra04]

• feat: preserve MCP tool annotations in bidirectional conversion --skip-pipeline #2746
• fix: prevent send on closed channel panic in provider queue shutdown --skip-pipeline #2725 👈 (View in Graphite)
• main

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

[coderabbitai]

📝 Walkthrough

Walkthrough

ProviderQueue no longer closes pq.queue; shutdown is driven by pq.signalClosing() / pq.done. Workers select on pq.queue and <-pq.done>, then drain remaining messages sending "provider is shutting down" errors. Update/Remove move or reroute buffered messages before signaling shutdown; context-cancel no longer returns pooled messages.

Changes

Cohort / File(s)	Summary
Core shutdown & worker logic core/bifrost.go	Removed closeQueue() and closeOnce; pq.queue remains open. Worker loop rewritten to select on pq.queue and <-pq.done>. Added drainQueueWithErrors. RemoveProvider signals closing, waits, then drains residuals before removing provider (provider-slice removal uses CAS). UpdateProvider transfers buffered messages to newPq before signalling old shutdown; overflow messages receive "queue full" errors; failures roll back by signalling and draining newPq. Producer paths (tryRequest/tryStreamRequest) attempt reroute to a non-closing replacement or return "provider is shutting down". On ctx.Done() they no longer call releaseChannelMessage(msg). Shutdown() signals all queues, waits for workers, then final-drains queues with errors.
Tests: concurrency & lifecycle core/bifrost_test.go	Added runtime and sync/atomic imports, test helpers (e.g., newTestChannelMessage) and many concurrency/lifecycle tests: deterministic TOCTOU race reproduction, isClosing()/signalClosing() invariants, worker shutdown without closing queue, buffered-drain error checks, UpdateProvider transfer/reroute/rollback tests, and RemoveProvider shutdown/TOCTOU stress tests ensuring no panics and correct error propagation.
Minor formatting transports/bifrost-http/handlers/providers.go	Removed two blank lines; no behavioral change.

Sequence Diagram(s)

sequenceDiagram
    participant Producer
    participant TryReq as tryRequest / tryStreamRequest
    participant ReqQueues as requestQueues
    participant OldPQ as ProviderQueue_old
    participant NewPQ as ProviderQueue_new
    participant Worker_old as Worker_old
    participant Done_old as pq.done_old
    participant MsgErr as msg.Err

    Producer->>TryReq: submit request
    TryReq->>ReqQueues: lookup provider queue
    ReqQueues->>OldPQ: return queue
    TryReq->>OldPQ: check isClosing()
    alt replacement exists
        TryReq->>NewPQ: reroute request
        NewPQ->>NewPQ: enqueue request
    else no replacement
        TryReq-->>Producer: "provider is shutting down"
    end

    Note over OldPQ,Done_old: Update/Remove transfer buffered items → then call pq.signalClosing()
    Caller->>OldPQ: pq.signalClosing()
    OldPQ->>Done_old: notify workers
    Worker_old->>OldPQ: drain residual messages
    loop per drained message
        Worker_old->>MsgErr: send "provider is shutting down"
    end
    Worker_old->>Worker_old: wg.Done() and exit

Loading

Estimated Code Review Effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Poem

🐇 I nudged the queues and hopped in line,

Moved buffered notes before the closing sign.
Workers listened, then gently sent their say,
“Provider is shutting down” — soft on the way.
The burrow hums; no panics chase the day.

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 93.75% which is sufficient. The required threshold is 80.00%.
Description check	✅ Passed	The PR description covers all required sections from the template with substantial detail on the problem, changes, testing approach, and checklist items.
Title check	✅ Passed	The PR title accurately describes the main change: preventing send-on-closed-channel panics during provider queue shutdown by fundamentally restructuring shutdown logic.

_{✏️ 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 04-15-fix_send_to_closed_channels_edge_case_race_condition_on_provider_reload

---

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

[github-actions]

🧪 Test Suite Available

This PR can be tested by a repository admin.

Run tests for PR #2725

[greptile-apps]

Confidence Score: 3/5

Not ready to merge — prior review comments flagged P1 reliability concerns (N-worker concurrent drain race, residual post-drain enqueue window) that remain open, and the PR is explicitly marked DO NOT MERGE.

The core channel-safety fix is sound and several prior concerns were addressed (double-release removed, time.After leak in drain loop removed, ctx.Done() added to stream response-wait). However, the concurrent N-worker drain race and the acknowledged post-drain TOCTOU window that can strand requests until context expiry are still open P1 concerns from earlier review rounds. New findings in this pass are P2 only (dead !ok test branches, inaccurate doc comment).

core/bifrost.go — the drain-loop concurrency and post-drain window; core/bifrost_test.go — dead !ok branches in two test goroutines

Important Files Changed

Filename	Overview
core/bifrost.go	Core fix: removes closeQueue() and switches workers to exit via pq.done instead of channel close. Adds drainQueueWithErrors(), re-routing logic in tryRequest/tryStreamRequest, ctx.Done() to tryStreamRequest response-wait, and rollback paths for UpdateProvider errors. A doc comment overstates the write-lock guarantee for stale producers.
core/bifrost_test.go	Adds comprehensive test suite covering the TOCTOU race, drain behaviour, worker lifecycle, and concurrent UpdateProvider/RemoveProvider scenarios. Two test goroutines retain a dead !ok check on pq.queue that contradicts the fix's never-close invariant.
transports/bifrost-http/handlers/providers.go	Trivial blank-line removals after function calls; no logic changes.

_{Reviews (12): Last reviewed commit: "fix: send to closed channels edge case r..." | Re-trigger Greptile}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/bifrost_test.go`:
- Around line 1304-1394: Test reproduces the old bug by calling pq.closeQueue()
directly; change TestProviderQueue_SendOnClosedChannel_Race to drive shutdown
via the real public flow (e.g. call RemoveProvider/UpdateProvider/Shutdown that
will invoke ProviderQueue.signalClosing()/closeQueue() internally) instead of
calling ProviderQueue.closeQueue() directly, and exercise the producer path
through tryRequest (or the code that sends to pq.queue with the pq.done guard)
so the test asserts there is no panic, no hang, and that the shutdown call
returns an error to the caller when appropriate; update assertions to fail on
any panic or hang and to validate the public API shutdown behavior rather than
expecting a panic from direct pq.closeQueue() use.

In `@core/bifrost.go`:
- Around line 4937-4971: After signalClosing(), producers that passed
isClosing() can still send into pq.queue and get stuck if workers exit on a
default empty-queue return; fix by adding a sender quiesce/refcount so workers
don't stop servicing the old queue until all pre-close senders complete.
Concretely: in tryRequest and tryStreamRequest increment a sender counter
(atomic or sync.WaitGroup) before attempting pq.queue <- msg and decrement after
the send (or on error/timeout), have signalClosing() wait for that counter to
reach zero (or call Done on the WaitGroup) before closing/marking pq.done, and
ensure queue closure is performed via sync.Once to avoid send-on-closed-channel
races; keep pq.queue reads draining until the refcount is zero and pq.done is
closed.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 268036c0-f15e-4768-9053-51df70324f75

📥 Commits

Reviewing files that changed from the base of the PR and between 67938e2 and 7b417a6.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (2)

core/bifrost.go (1)

4955-4981: ⚠️ Potential issue | 🔴 Critical

Workers can still exit while stale producers enqueue into the old queue.

After done is closed, stale producers can still win pq.queue <- msg over <-pq.done, and this drain loop can return on a transient empty/default before those sends occur. That leaves messages in an open queue with no workers, causing caller hangs until context timeout.

Suggested fix (sender quiesce barrier)

 type ProviderQueue struct {
 	queue      chan *ChannelMessage
 	done       chan struct{}
 	closing    uint32
+	activeSenders atomic.Int64
 	signalOnce sync.Once
 	closeOnce  sync.Once
 }
+
+func (pq *ProviderQueue) beginSend() bool {
+	if pq.isClosing() {
+		return false
+	}
+	pq.activeSenders.Add(1)
+	if pq.isClosing() {
+		pq.activeSenders.Add(-1)
+		return false
+	}
+	return true
+}
+
+func (pq *ProviderQueue) endSend() {
+	pq.activeSenders.Add(-1)
+}

- select {
+ if !pq.beginSend() {
+   bifrost.releaseChannelMessage(msg)
+   return nil, newBifrostErrorFromMsg("provider is shutting down")
+ }
+ defer pq.endSend()
+ select {
   case pq.queue <- msg:
   case <-pq.done:

-               default:
-                   return
+               default:
+                   if pq.activeSenders.Load() == 0 {
+                       return
+                   }
+                   time.Sleep(1 * time.Millisecond)

# Stress check for orphaned request hangs under shutdown/update races
go test -run TestProviderQueue_SendOnClosedChannel_Race ./core -count=200 -race -timeout=60s

As per coding guidelines, "Use ProviderQueue pattern with atomic flags and sync.Once to prevent 'send on closed channel' panics when managing provider request queues."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 4955 - 4981, The drain loop can return while
producers still can succeed in sending to pq.queue after pq.done is closed,
leaving messages orphaned; fix by implementing a quiesce/close protocol on
ProviderQueue: add an atomic boolean (e.g. pq.closed) and a sync.Once close
operation that marks pq.closed and prevents further sends, make the enqueue path
check pq.closed (return an error rather than blocking) before attempting
pq.queue <- msg, and have shutdown use the sync.Once to close and then wait for
workers to fully drain the queue (or use a WaitGroup) before returning; update
the drain loop (the section reading from pq.queue and using
releaseChannelMessage) to rely on the closed flag + empty queue semantics so no
messages can be enqueued after shutdown begins.

core/bifrost_test.go (1)

1306-1396: ⚠️ Potential issue | 🟠 Major

This test still validates the old panic path, not the fixed shutdown contract.

Line 1381 manually calls pq.closeQueue(), and Lines 1392-1394 require panic occurrence for success. That makes CI success depend on a deprecated unsafe path rather than RemoveProvider/UpdateProvider/Shutdown behavior after the fix. Please remove or rewrite this test to drive the public shutdown/update flow and assert no panic, no hang, and shutdown error propagation.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost_test.go` around lines 1306 - 1396, The test
TestProviderQueue_SendOnClosedChannel_Race still forces the old unsafe path by
calling pq.closeQueue() directly and expecting a panic; instead rewrite it to
exercise the public shutdown/update API (e.g. call the real
RemoveProvider/UpdateProvider or Shutdown code paths that replace
pq.signalClosing()/pq.closeQueue()) and assert the new contract: no "send on
closed channel" panic, no hangs, and proper error propagation from the public
shutdown/update calls. Replace direct calls to
pq.closeQueue()/pq.signalClosing() with the higher-level methods that operate on
ProviderQueue (the code paths that tryRequest uses: isClosing(), tryRequest, and
the public RemoveProvider/UpdateProvider/Shutdown flows), run many iterations to
exercise races but fail the test if any panic/hang occurs, and add assertions
that the public call returns the expected shutdown/update error where
applicable.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/bifrost.go`:
- Around line 4386-4388: The code currently calls bifrost.getProviderQueue in
the closing-path reroute which lazily creates a provider queue and can resurrect
a provider during shutdown; change the lookup to a non-creating accessor (e.g. a
getProviderQueueNoCreate / lookupProviderQueue / providerQueueIfExists method)
so it only returns an existing queue or an error without creating one, update
the branch that assigns pq to use that non-creating lookup instead of
bifrost.getProviderQueue, and make the same replacement at the other occurrence
mentioned (around lines 4638-4640) and in RemoveProvider-related cleanup paths
to avoid recreating providers during shutdown.

---

Duplicate comments:
In `@core/bifrost_test.go`:
- Around line 1306-1396: The test TestProviderQueue_SendOnClosedChannel_Race
still forces the old unsafe path by calling pq.closeQueue() directly and
expecting a panic; instead rewrite it to exercise the public shutdown/update API
(e.g. call the real RemoveProvider/UpdateProvider or Shutdown code paths that
replace pq.signalClosing()/pq.closeQueue()) and assert the new contract: no
"send on closed channel" panic, no hangs, and proper error propagation from the
public shutdown/update calls. Replace direct calls to
pq.closeQueue()/pq.signalClosing() with the higher-level methods that operate on
ProviderQueue (the code paths that tryRequest uses: isClosing(), tryRequest, and
the public RemoveProvider/UpdateProvider/Shutdown flows), run many iterations to
exercise races but fail the test if any panic/hang occurs, and add assertions
that the public call returns the expected shutdown/update error where
applicable.

In `@core/bifrost.go`:
- Around line 4955-4981: The drain loop can return while producers still can
succeed in sending to pq.queue after pq.done is closed, leaving messages
orphaned; fix by implementing a quiesce/close protocol on ProviderQueue: add an
atomic boolean (e.g. pq.closed) and a sync.Once close operation that marks
pq.closed and prevents further sends, make the enqueue path check pq.closed
(return an error rather than blocking) before attempting pq.queue <- msg, and
have shutdown use the sync.Once to close and then wait for workers to fully
drain the queue (or use a WaitGroup) before returning; update the drain loop
(the section reading from pq.queue and using releaseChannelMessage) to rely on
the closed flag + empty queue semantics so no messages can be enqueued after
shutdown begins.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 36a8b354-7f23-48c7-bc89-0cf443399215

📥 Commits

Reviewing files that changed from the base of the PR and between 7b417a6 and 592c6e1.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (1)

core/bifrost.go (1)

6069-6108: ⚠️ Potential issue | 🔴 Critical

The one-shot final drain still leaves hung requests possible.

drainQueueWithErrors is still a non-blocking sweep after the workers exit. A producer that already holds pq, races past the last default: return, and then wins pq.queue <- msg lands a message in an open queue with no workers and no further drain pass, so the caller blocks on msg.Response / msg.Err until its context is cancelled. This narrows the old panic window, but it doesn't actually close the shutdown race. As per coding guidelines, "Use ProviderQueue pattern with atomic flags and sync.Once to prevent 'send on closed channel' panics when managing provider request queues."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 6069 - 6108, drainQueueWithErrors currently
does a one-shot non-blocking sweep and can miss messages sent after it returns;
fix by adding a ProviderQueue-level closed flag (atomic.Bool) plus a sync.Once
close method and make producers check that flag before attempting pq.queue <-
msg (e.g. in tryRequest or the send path) so no new sends are allowed after
close is initiated; call the closeOnce/atomic store at the start of
drainQueueWithErrors (or in the same shutdown path that waits for workers) to
atomically prevent further sends to pq.queue, then keep the existing draining
logic to flush existing messages and return. Use the existing symbols:
ProviderQueue, pq.queue, drainQueueWithErrors, and the producer send site
(tryRequest) to implement the atomic closed flag and sync.Once close.

🧹 Nitpick comments (1)

core/bifrost_test.go (1)

1976-1998: Ordering differs from TestUpdateProvider_TransferOrdering — clarify step comments.

The updater loop here performs:

1. signalClosing(oldPq) (line 1977)
2. Drain oldPq → newPq (lines 1982-1995)
3. Store newPq (line 1998)

But TestUpdateProvider_TransferOrdering (lines 1814-1838) expects the inverse: transfer before signalClosing().

For this test's goal (zero panics), either ordering is safe since the queue channel is never closed. However, the comment "performs UpdateProvider-style queue replacements" is misleading if the production UpdateProvider actually transfers before signaling.

Consider updating the step comments to clarify this is an alternative ordering for stress-testing purposes, or align with the expected production ordering:

-// Step 2: signal old queue closing so no new producers land on it.
-oldPq.signalClosing()
-
-// Step 3: drain oldPq into newPq before exposing newPq...
+// Note: For this stress test, we signal closing first so producers
+// see isClosing() == true and re-route. Production UpdateProvider
+// may use different ordering (transfer-then-signal).
+oldPq.signalClosing()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost_test.go` around lines 1976 - 1998, The test's updater loop
currently calls oldPq.signalClosing() before draining oldPq into newPq and then
does requestQueues.Store(provider, newPq), which conflicts with the ordering
assumed by TestUpdateProvider_TransferOrdering (which transfers before
signaling); either update the inline step comments to explicitly state this test
intentionally uses the alternative ordering for stress-testing (mentioning
oldPq.signalClosing(), the drain loop, and requestQueues.Store(provider,
newPq)), or change the sequence to match production by moving the
signalClosing() call to after the drain so the flow matches
TestUpdateProvider_TransferOrdering; keep the comment clear about which ordering
is being tested.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/bifrost.go`:
- Around line 3238-3240: Currently requestQueues.Store(providerKey, newPq)
publishes newPq before any workers are started, causing the transfer goroutine
to block if there are more messages than newPq can accept; fix by keeping newPq
local (do not call requestQueues.Store) until after createBaseProvider and the
new provider/workers are fully initialized and transferWaitGroup has been
satisfied, start the transfer goroutine only after workers are consuming
newPq.queue, and if any initialization or the provider-slice swap
(createBaseProvider, provider slice update) fails roll back or avoid storing the
entry (or remove it) so callers cannot enqueue into a visible queue with zero
workers; apply the same change to the other affected blocks around the transfer
goroutine sections noted (including the regions around lines 3252-3294 and
3317-3324).
- Around line 4772-4774: The code currently returns pooled ChannelMessage to the
pool on ctx.Done() by calling bifrost.releaseChannelMessage(msg), which can let
a late write to req.ResponseStream or req.Err leak stale data into the next
user; remove the call to bifrost.releaseChannelMessage(msg) from the case
<-ctx.Done() branch so the worker retains ownership of the ChannelMessage until
it finishes handling the request and explicitly calls
bifrost.releaseChannelMessage(msg) after it sends to req.ResponseStream or
req.Err; keep the newBifrostCtxDoneError return but ensure any code paths that
complete the worker (the sender to req.ResponseStream/req.Err) perform the
pool.Put/reset via releaseChannelMessage to avoid stale fields.

---

Duplicate comments:
In `@core/bifrost.go`:
- Around line 6069-6108: drainQueueWithErrors currently does a one-shot
non-blocking sweep and can miss messages sent after it returns; fix by adding a
ProviderQueue-level closed flag (atomic.Bool) plus a sync.Once close method and
make producers check that flag before attempting pq.queue <- msg (e.g. in
tryRequest or the send path) so no new sends are allowed after close is
initiated; call the closeOnce/atomic store at the start of drainQueueWithErrors
(or in the same shutdown path that waits for workers) to atomically prevent
further sends to pq.queue, then keep the existing draining logic to flush
existing messages and return. Use the existing symbols: ProviderQueue, pq.queue,
drainQueueWithErrors, and the producer send site (tryRequest) to implement the
atomic closed flag and sync.Once close.

---

Nitpick comments:
In `@core/bifrost_test.go`:
- Around line 1976-1998: The test's updater loop currently calls
oldPq.signalClosing() before draining oldPq into newPq and then does
requestQueues.Store(provider, newPq), which conflicts with the ordering assumed
by TestUpdateProvider_TransferOrdering (which transfers before signaling);
either update the inline step comments to explicitly state this test
intentionally uses the alternative ordering for stress-testing (mentioning
oldPq.signalClosing(), the drain loop, and requestQueues.Store(provider,
newPq)), or change the sequence to match production by moving the
signalClosing() call to after the drain so the flow matches
TestUpdateProvider_TransferOrdering; keep the comment clear about which ordering
is being tested.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: e081032f-40d5-4ef9-92cd-58be67d2a9bd

📥 Commits

Reviewing files that changed from the base of the PR and between 7993b5e and 225802b.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (1)

core/bifrost.go (1)

3145-3146: ⚠️ Potential issue | 🟠 Major

The final sweep still leaves a stranded-request window.

The helper comment at Line 6100 describes the remaining race: a producer that passed the closing gate before signalClosing() can still win pq.queue <- msg after this one-shot drain exits via default. At that point RemoveProvider, UpdateProvider, and Shutdown have no worker and no second sweep for that queue, so the caller waits on msg.Response/msg.Err until its own context is cancelled.

This fixes the panic, but shutdown/update still isn't deterministic under load. You'll need a sender quiesce/barrier or an equivalent always-on drainer before treating the queue as fully drained.

Also applies to: 3321-3322, 6100-6131, 6656-6660

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 3145 - 3146, The final one-shot drain in
drainQueueWithErrors leaves a race where a producer that passed signalClosing()
can still send to pq.queue (pq.queue <- msg) after the default-exit drain,
leaving callers blocked on msg.Response/msg.Err; fix by introducing a sender
quiesce/barrier or an always-on drainer: add a quiesce phase or a
"no-new-senders" signal that producers must observe (or a producer WaitGroup)
before treating the queue drained, and replace the one-shot default-exit with a
persistent goroutine that continues draining pq.queue until a proper
closed+empty condition is reached; update signalClosing(), drainQueueWithErrors,
and the shutdown paths for RemoveProvider/UpdateProvider/Shutdown to coordinate
with the quiesce (or persistent drainer) so no msg can be sent after the closing
gate without a worker to complete msg.Response/msg.Err.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/bifrost.go`:
- Around line 3238-3240: The new per-provider queue (newPq) is stored into
bifrost.requestQueues (bifrost.requestQueues.Store(providerKey, newPq)) before
the CAS replacement loop can fail, leaving a workerless queue if
maxReplacementAttempts is reached; modify the failure path in the CAS loop (the
branch that gives up after maxReplacementAttempts) to undo the store by
removing/clearing the entry for providerKey (i.e., rollback newPq) so callers
cannot enqueue into a queue with no workers, and apply the same rollback logic
to the analogous failure path around the second occurrence (the block referenced
at lines ~3345-3351) to ensure both replacement code paths remove newPq when
they give up.

---

Duplicate comments:
In `@core/bifrost.go`:
- Around line 3145-3146: The final one-shot drain in drainQueueWithErrors leaves
a race where a producer that passed signalClosing() can still send to pq.queue
(pq.queue <- msg) after the default-exit drain, leaving callers blocked on
msg.Response/msg.Err; fix by introducing a sender quiesce/barrier or an
always-on drainer: add a quiesce phase or a "no-new-senders" signal that
producers must observe (or a producer WaitGroup) before treating the queue
drained, and replace the one-shot default-exit with a persistent goroutine that
continues draining pq.queue until a proper closed+empty condition is reached;
update signalClosing(), drainQueueWithErrors, and the shutdown paths for
RemoveProvider/UpdateProvider/Shutdown to coordinate with the quiesce (or
persistent drainer) so no msg can be sent after the closing gate without a
worker to complete msg.Response/msg.Err.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 21f8ac1d-5408-4c3a-8c3f-8274b9e756e9

📥 Commits

Reviewing files that changed from the base of the PR and between d0f3b1c and 24896d4.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

🚧 Files skipped from review as they are similar to previous changes (1)

• core/bifrost_test.go

[coderabbitai]

🧹 Nitpick comments (1)

core/bifrost_test.go (1)

1333-1337: Consider extracting a shared newTestProviderQueue helper to reduce drift.

ProviderQueue test setup is repeated many times with the same fields. Centralizing it will make future queue-lifecycle changes less error-prone.

♻️ Suggested refactor

+func newTestProviderQueue(bufferSize int) *ProviderQueue {
+	return &ProviderQueue{
+		queue:      make(chan *ChannelMessage, bufferSize),
+		done:       make(chan struct{}),
+		signalOnce: sync.Once{},
+	}
+}

- pq := &ProviderQueue{
- 	queue:      make(chan *ChannelMessage, 10),
- 	done:       make(chan struct{}),
- 	signalOnce: sync.Once{},
- }
+ pq := newTestProviderQueue(10)

Also applies to: 1429-1433, 1476-1480, 1502-1506, 1555-1559, 1636-1640, 1721-1730, 1793-1802, 1875-1879, 2035-2039, 2095-2099, 2166-2170, 2226-2230

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost_test.go` around lines 1333 - 1337, Several tests repeatedly
instantiate ProviderQueue with identical fields; extract a helper function named
newTestProviderQueue that returns *ProviderQueue configured with queue:
make(chan *ChannelMessage, 10), done: make(chan struct{}), and signalOnce:
sync.Once{} to centralize setup and reduce duplication; replace inline struct
literals in tests (the places creating &ProviderQueue{queue: make(chan
*ChannelMessage, 10), done: make(chan struct{}), signalOnce: sync.Once{}}) with
calls to newTestProviderQueue() so future lifecycle changes are made in one
location.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@core/bifrost_test.go`:
- Around line 1333-1337: Several tests repeatedly instantiate ProviderQueue with
identical fields; extract a helper function named newTestProviderQueue that
returns *ProviderQueue configured with queue: make(chan *ChannelMessage, 10),
done: make(chan struct{}), and signalOnce: sync.Once{} to centralize setup and
reduce duplication; replace inline struct literals in tests (the places creating
&ProviderQueue{queue: make(chan *ChannelMessage, 10), done: make(chan struct{}),
signalOnce: sync.Once{}}) with calls to newTestProviderQueue() so future
lifecycle changes are made in one location.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 3e5e3419-e0a0-4658-a355-73bd91a132ae

📥 Commits

Reviewing files that changed from the base of the PR and between 24896d4 and 5018b50.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

🚧 Files skipped from review as they are similar to previous changes (1)

• core/bifrost.go

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

core/bifrost.go (1)

6135-6158: ⚠️ Potential issue | 🟠 Major

Reset the full pooled ChannelMessage before Put().

releaseChannelMessage() returns the object to channelMessagePool, but it keeps Context and the embedded BifrostRequest populated. That retains old request payloads/contexts in the pool and violates the repo’s pool-reset rule.

♻️ Suggested fix

 func (bifrost *Bifrost) releaseChannelMessage(msg *ChannelMessage) {
 	// Put channels back in pools
 	bifrost.responseChannelPool.Put(msg.Response)
 	bifrost.errorChannelPool.Put(msg.Err)
@@
 	// Clear references and return to pool
+	msg.BifrostRequest = schemas.BifrostRequest{}
+	msg.Context = nil
 	msg.Response = nil
 	msg.ResponseStream = nil
 	msg.Err = nil
 	bifrost.channelMessagePool.Put(msg)
 }

As per coding guidelines, "Always reset ALL fields of pooled objects before calling pool.Put() to prevent stale data from leaking to subsequent users."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 6135 - 6158, The releaseChannelMessage function
returns a ChannelMessage to channelMessagePool but currently leaves msg.Context
and the embedded BifrostRequest (and any other request-related fields)
populated, which can leak stale data; before calling
bifrost.channelMessagePool.Put(msg) clear all remaining fields on msg — at
minimum set msg.Context = nil and zero out the embedded BifrostRequest (and any
other request/payload fields) after draining/returning channels and before Put —
so the pooled ChannelMessage contains no prior request state when reused.

♻️ Duplicate comments (1)

core/bifrost.go (1)

6094-6133: ⚠️ Potential issue | 🟠 Major

This still leaves a post-drain enqueue window.

The comment here is accurate, but it also describes a real unresolved bug: after the worker drain loop and this final sweep both exit on default, a producer that already passed the closing gate can still have select choose pq.queue <- msg over <-pq.done. That request has no worker and no later sweep, so it hangs until its own context is cancelled.

This is the same stale-producer race previously called out; the single extra sweep narrows it but does not close it. A sender quiesce/barrier is still needed before the last drain returns.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 6094 - 6133, The final non-blocking sweep in
drainQueueWithErrors still leaves a race where a producer that passed the
closing gate can enqueue after the sweep and never get drained; fix by adding a
sender-side quiesce: add a sender counter/WG to ProviderQueue (e.g., senderWg or
activeSenders int32 + sync.WaitGroup) and update tryRequest to increment the
sender ref before checking/attempting the send and decrement after the send or
on context cancel; then change drainQueueWithErrors to wait for that senderWg to
reach zero before doing the final non-blocking sweep so no sender can enqueue
after the sweep exits. Ensure updates touch ProviderQueue struct, tryRequest
(the producer path), and drainQueueWithErrors to use the new sender quiesce
mechanism.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@core/bifrost.go`:
- Around line 3325-3334: The rollback path currently removes requestQueues and
waitGroups but leaves the stale provider in bifrost.providers, causing
GetConfiguredProviders() and getProviderByKey() to expose a dead provider and
allowing prepareProvider() to append a duplicate; update the failure branches
that call newPq.signalClosing(), bifrost.requestQueues.Delete(providerKey),
bifrost.waitGroups.Delete(providerKey), and bifrost.drainQueueWithErrors(newPq)
to also remove the provider entry from bifrost.providers (and any other
provider-index maps) for the same providerKey so the provider is fully removed
on error; apply the same change to both failure locations (the block around
drainQueueWithErrors and the similar block at 3346-3353) to ensure retries via
prepareProvider() start from a clean state.

---

Outside diff comments:
In `@core/bifrost.go`:
- Around line 6135-6158: The releaseChannelMessage function returns a
ChannelMessage to channelMessagePool but currently leaves msg.Context and the
embedded BifrostRequest (and any other request-related fields) populated, which
can leak stale data; before calling bifrost.channelMessagePool.Put(msg) clear
all remaining fields on msg — at minimum set msg.Context = nil and zero out the
embedded BifrostRequest (and any other request/payload fields) after
draining/returning channels and before Put — so the pooled ChannelMessage
contains no prior request state when reused.

---

Duplicate comments:
In `@core/bifrost.go`:
- Around line 6094-6133: The final non-blocking sweep in drainQueueWithErrors
still leaves a race where a producer that passed the closing gate can enqueue
after the sweep and never get drained; fix by adding a sender-side quiesce: add
a sender counter/WG to ProviderQueue (e.g., senderWg or activeSenders int32 +
sync.WaitGroup) and update tryRequest to increment the sender ref before
checking/attempting the send and decrement after the send or on context cancel;
then change drainQueueWithErrors to wait for that senderWg to reach zero before
doing the final non-blocking sweep so no sender can enqueue after the sweep
exits. Ensure updates touch ProviderQueue struct, tryRequest (the producer
path), and drainQueueWithErrors to use the new sender quiesce mechanism.

🪄 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: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ab83018b-f961-44e3-8992-e5bd265ba4c9

📥 Commits

Reviewing files that changed from the base of the PR and between 5018b50 and 99c287f.

📒 Files selected for processing (2)

• core/bifrost.go
• core/bifrost_test.go

🚧 Files skipped from review as they are similar to previous changes (1)

• core/bifrost_test.go

[coderabbitai]

♻️ Duplicate comments (1)

core/bifrost.go (1)

6136-6142: ⚠️ Potential issue | 🟠 Major

Residual post-sweep enqueue can still strand requests.

This one-shot non-blocking sweep still allows a producer to enqueue after the final default exit, leaving the request without workers and waiting until caller cancellation. That’s a reliability gap during RemoveProvider/UpdateProvider/Shutdown when contexts are long-lived.

Please add sender quiescing (refcount/barrier) before final drain, or keep draining until producer activity is fully quiesced.

As per coding guidelines, "Use ProviderQueue pattern with atomic flags and sync.Once to prevent 'send on closed channel' panics when managing provider request queues."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 6136 - 6142, The one-shot non-blocking sweep
can leave producers able to enqueue after the final `select { default: }` exit,
stranding requests; update the drain logic in the provider queue code (the sweep
invoked during `RemoveProvider`/`UpdateProvider`/`Shutdown`) to wait for sender
quiescence instead of a single non-blocking pass: add a sender-side reference
count/barrier (e.g., increment in send path and decrement on send completion)
and have the drain loop wait until that refcount reaches zero or a short
timeout, and protect queue closure with a `ProviderQueue` pattern using an
atomic closed flag and `sync.Once` to ensure no sends occur to a closed channel
and avoid race panics. Ensure you reference and modify the sweep/drain routine,
the send path that currently does the non-blocking `select { default: }`, and
the provider lifecycle callers (`RemoveProvider`, `UpdateProvider`, `Shutdown`)
so they trigger quiescing before the final drain.

🧹 Nitpick comments (2)

core/bifrost_test.go (1)

2223-2299: Rename this test or assert the drain outcome.

The body currently treats both successfulSends and shutdownErrors as valid, so it proves “no panic in the TOCTOU window” rather than “new producers get shutdown errors.” Renaming it to reflect that narrower guarantee, or draining/asserting the successfully enqueued messages, would make the test intent match its behavior.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost_test.go` around lines 2223 - 2299, The test
TestRemoveProvider_ConcurrentNewProducersDuringShutdown currently allows both
successfulSends and shutdownErrors, so update it to assert the drain outcome:
after producerWg.Wait() drain pq.queue non-blockingly (read until channel empty)
counting drainedMessages, then assert atomic.LoadInt64(&successfulSends) ==
int64(drainedMessages) and optionally assert drainedMessages == 0 if the intent
is "new producers get shutdown errors"; reference
TestRemoveProvider_ConcurrentNewProducersDuringShutdown, pq.queue,
successfulSends and shutdownErrors when making the change.

core/bifrost.go (1)

3186-3192: Concurrency guarantee comment is currently stronger than the implementation.

Line 3187/3188 says all producers go through getProviderQueue read-locking, but stale-producer reroute paths use direct requestQueues.Load (Line 4465, Line 4732). Please reword this comment to avoid a misleading guarantee.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@core/bifrost.go` around lines 3186 - 3192, The comment on UpdateProvider
overstates the concurrency guarantee; update it to say that while producer paths
that call getProviderQueue do acquire the providerMutex read lock before looking
up/enqueueing (so they cannot see newPq until the write lock is released and
workers are started), there are other paths (notably the stale-producer reroute
code which performs direct requestQueues.Load) that bypass getProviderQueue and
its read-lock, so UpdateProvider should not claim an absolute prohibition of any
producer observing newPq before workers run; mention getProviderQueue and the
direct requestQueues.Load/stale-producer reroute paths by name to clarify the
exception.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@core/bifrost.go`:
- Around line 6136-6142: The one-shot non-blocking sweep can leave producers
able to enqueue after the final `select { default: }` exit, stranding requests;
update the drain logic in the provider queue code (the sweep invoked during
`RemoveProvider`/`UpdateProvider`/`Shutdown`) to wait for sender quiescence
instead of a single non-blocking pass: add a sender-side reference count/barrier
(e.g., increment in send path and decrement on send completion) and have the
drain loop wait until that refcount reaches zero or a short timeout, and protect
queue closure with a `ProviderQueue` pattern using an atomic closed flag and
`sync.Once` to ensure no sends occur to a closed channel and avoid race panics.
Ensure you reference and modify the sweep/drain routine, the send path that
currently does the non-blocking `select { default: }`, and the provider
lifecycle callers (`RemoveProvider`, `UpdateProvider`, `Shutdown`) so they
trigger quiescing before the final drain.

---

Nitpick comments:
In `@core/bifrost_test.go`:
- Around line 2223-2299: The test
TestRemoveProvider_ConcurrentNewProducersDuringShutdown currently allows both
successfulSends and shutdownErrors, so update it to assert the drain outcome:
after producerWg.Wait() drain pq.queue non-blockingly (read until channel empty)
counting drainedMessages, then assert atomic.LoadInt64(&successfulSends) ==
int64(drainedMessages) and optionally assert drainedMessages == 0 if the intent
is "new producers get shutdown errors"; reference
TestRemoveProvider_ConcurrentNewProducersDuringShutdown, pq.queue,
successfulSends and shutdownErrors when making the change.

In `@core/bifrost.go`:
- Around line 3186-3192: The comment on UpdateProvider overstates the
concurrency guarantee; update it to say that while producer paths that call
getProviderQueue do acquire the providerMutex read lock before looking
up/enqueueing (so they cannot see newPq until the write lock is released and
workers are started), there are other paths (notably the stale-producer reroute
code which performs direct requestQueues.Load) that bypass getProviderQueue and
its read-lock, so UpdateProvider should not claim an absolute prohibition of any
producer observing newPq before workers run; mention getProviderQueue and the
direct requestQueues.Load/stale-producer reroute paths by name to clarify the
exception.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: c1ae4c33-0ca0-4499-b119-35c10d9b3bae

📥 Commits

Reviewing files that changed from the base of the PR and between 99c287f and 8e2fab7.

📒 Files selected for processing (3)

• core/bifrost.go
• core/bifrost_test.go
• transports/bifrost-http/handlers/providers.go

💤 Files with no reviewable changes (1)

• transports/bifrost-http/handlers/providers.go

[Pratham-Mishra04]

@greptile

Re: N-worker concurrent drain race

Go channels guarantee that each message is received by exactly one goroutine — this is a language-level invariant, not
a convention. When all N workers select on pq.done, whichever worker wins the pq.done case is responsible for draining;
the others move on. Even if multiple workers simultaneously enter the drain loop, each <-pq.queue dequeue is atomic
and exclusive — there is no double-receive. No shared mutable state is touched during the drain read itself, so
concurrent drainers are safe.

Re: post-drain TOCTOU window

This is a known, documented limitation. The window exists between a producer passing the isClosing() check and
depositing its message, after the drain sweep has already exited. The message will sit unread until the caller's
context deadline fires, at which point the ctx.Done() case in the response-wait select returns an error to the caller.
The window is therefore bounded by the request's own deadline — no request is stranded indefinitely. Eliminating it
entirely would require sender-side reference counting on the queue (increment before the isClosing() check, decrement
after enqueue, drain waits for zero) which adds overhead on every single enqueue in the hot path. We've judged the
bounded-by-deadline behavior acceptable for now and left the tradeoff documented in the drainQueueWithErrors comment.

[greptile-apps]

Both points are well-reasoned. My original concerns are addressed:

On concurrent drain: You're correct — channel receives in Go are atomic and mutually exclusive by the memory model. Multiple goroutines racing on <-pq.queue will each get distinct messages, never the same one. The "concurrent drainers" concern I raised was unfounded. Withdrawing that.

On post-drain TOCTOU: The bounded-by-deadline framing is the right mental model here. A message that slips through after the drain sweep isn't lost forever — it waits at most until ctx.Done() fires on the producer side. The tradeoff you've described (sender-side WaitGroup / refcount on the hot enqueue path) is real overhead that isn't justified when the deadline already provides the bound. The comment in drainQueueWithErrors documenting this tradeoff is the right call.

Both points accepted. LGTM on the concurrency reasoning.

Tip: You can give me more context with [.greptile/rules](https://www.greptile.com/docs/code-review/custom-standards) so I get this right next time.

[akshaydeo]

Merge activity

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