](https://app.getpostman.com/run-collection/31642484-2ba0e658-4dcd-49f4-845a-0c7ed745b916?action=collection%2Ffork&source=rip_markdown&collection-url=entityId%3D31642484-2ba0e658-4dcd-49f4-845a-0c7ed745b916%26entityType%3Dcollection%26workspaceId%3D63e853c8-9aec-477f-909c-7f02f543150e)
diff --git a/cli/go.mod b/cli/go.mod
index c887aac505..f7260a3b86 100644
--- a/cli/go.mod
+++ b/cli/go.mod
@@ -1,6 +1,6 @@
module github.com/maximhq/bifrost/cli
-go 1.26.2
+go 1.26.1
require (
github.com/bytedance/sonic v1.15.0
diff --git a/core/bifrost.go b/core/bifrost.go
index 7db6663758..70454a96e6 100644
--- a/core/bifrost.go
+++ b/core/bifrost.go
@@ -89,12 +89,41 @@ type Bifrost struct {
// ProviderQueue wraps a provider's request channel with lifecycle management
// to prevent "send on closed channel" panics during provider removal/update.
// Producers must check the closing flag or select on the done channel before sending.
+//
+// Why pq.queue is NEVER closed:
+//
+// Closing a channel in Go causes any concurrent send to that channel to panic
+// ("send on closed channel"). There is always a TOCTOU window between a
+// producer's isClosing() check and its select { case pq.queue <- msg: ... }:
+// the producer could pass isClosing() while the queue is open, get preempted,
+// and resume only after the queue is closed. Go's selectgo evaluates select
+// cases in a random order, so even having case <-pq.done: in the same select
+// does not protect against this — if selectgo evaluates the send case first on
+// a closed channel it panics immediately via goto sclose, before reaching done.
+//
+// To close pq.queue safely you would need a sender-side WaitGroup so that
+// signalClosing could wait for every in-flight producer to finish. That adds
+// non-trivial overhead on the hot request path.
+//
+// Instead, pq.done is the sole shutdown signal. Receiving from a closed channel
+// is always safe (returns the zero value immediately), so:
+// - Workers exit via case <-pq.done: — safe
+// - Producers bail via case <-pq.done: — safe
+// - drainQueueWithErrors handles any messages that slip through the TOCTOU window
+//
+// pq.queue is garbage collected automatically:
+// - RemoveProvider calls requestQueues.Delete, dropping the map's reference.
+// - UpdateProvider calls requestQueues.Store with a new queue, dropping the
+// map's reference to oldPq. Shutdown does not Delete at all — the whole
+// Bifrost instance is torn down.
+// In all cases, once no producer goroutine holds a reference to the
+// ProviderQueue, both the struct and pq.queue are eligible for GC.
+// No explicit close is needed.
type ProviderQueue struct {
- queue chan *ChannelMessage // the actual request queue channel
- done chan struct{} // closed to signal shutdown to producers
+ queue chan *ChannelMessage // the actual request queue channel — never closed, see above
+ done chan struct{} // closed by signalClosing() to signal shutdown; never written to otherwise
closing uint32 // atomic: 0 = open, 1 = closing
signalOnce sync.Once
- closeOnce sync.Once
}
func isLargePayloadPassthrough(ctx *schemas.BifrostContext) bool {
@@ -122,14 +151,6 @@ func (pq *ProviderQueue) signalClosing() {
})
}
-// closeQueue closes the provider queue.
-// Protected by sync.Once to prevent double-close.
-func (pq *ProviderQueue) closeQueue() {
- pq.closeOnce.Do(func() {
- close(pq.queue)
- })
-}
-
// isClosing returns true if the provider queue is closing.
// Uses atomic load for lock-free checking.
func (pq *ProviderQueue) isClosing() bool {
@@ -3109,57 +3130,36 @@ func (bifrost *Bifrost) RemoveProvider(providerKey schemas.ModelProvider) error
}
pq := pqValue.(*ProviderQueue)
- // Step 2: Signal closing to producers (prevents new sends)
- // This must happen before closing the queue to avoid "send on closed channel" panics
+ // Step 2: Signal closing. Blocks new producers (isClosing() returns true) and
+ // causes idle workers to drain remaining buffered requests with errors then exit.
pq.signalClosing()
bifrost.logger.Debug("signaled closing for provider %s", providerKey)
- // Step 3: Now safe to close the queue (no new producers can send)
- pq.closeQueue()
- bifrost.logger.Debug("closed request queue for provider %s", providerKey)
-
- // Step 4: Wait for all workers to finish processing in-flight requests
+ // Step 3: Wait for all workers to finish in-flight requests and exit.
waitGroup, exists := bifrost.waitGroups.Load(providerKey)
if exists {
waitGroup.(*sync.WaitGroup).Wait()
bifrost.logger.Debug("all workers for provider %s have stopped", providerKey)
}
- // Step 5: Remove the provider from the request queues
+ // Step 3b: Final drain sweep — see drainQueueWithErrors for full explanation.
+ bifrost.drainQueueWithErrors(pq)
+
+ // Step 4: Remove the provider from the request queues.
bifrost.requestQueues.Delete(providerKey)
- // Step 6: Remove the provider from the wait groups
+ // Step 5: Remove the provider from the wait groups.
bifrost.waitGroups.Delete(providerKey)
- // Step 7: Remove the provider from the providers slice
- replacementAttempts := 0
- maxReplacementAttempts := 100 // Prevent infinite loops in high-contention scenarios
- for {
- replacementAttempts++
- if replacementAttempts > maxReplacementAttempts {
- return fmt.Errorf("failed to replace provider %s in providers slice after %d attempts", providerKey, maxReplacementAttempts)
- }
- oldPtr := bifrost.providers.Load()
- var oldSlice []schemas.Provider
- if oldPtr != nil {
- oldSlice = *oldPtr
- }
- // Create new slice without the old provider of this key
- // Use exact capacity to avoid allocations
- if len(oldSlice) == 0 {
- return fmt.Errorf("provider %s not found in providers slice", providerKey)
- }
- newSlice := make([]schemas.Provider, 0, len(oldSlice)-1)
- for _, existingProvider := range oldSlice {
- if existingProvider.GetProviderKey() != providerKey {
- newSlice = append(newSlice, existingProvider)
- }
- }
- if bifrost.providers.CompareAndSwap(oldPtr, &newSlice) {
- bifrost.logger.Debug("successfully removed provider instance for %s in providers slice", providerKey)
- break
- }
- // Retrying as swapping did not work (likely due to concurrent modification)
+ // Step 6: Remove the provider from the providers slice.
+ if err := bifrost.removeProviderFromSlice(providerKey); err != nil {
+ bifrost.logger.Error(
+ "provider %s was removed from queues but could not be removed from the providers slice — "+
+ "bifrost.providers is now inconsistent. "+
+ "To recover: retry RemoveProvider(%s), or restart Bifrost if that fails.",
+ providerKey, providerKey,
+ )
+ return err
}
bifrost.logger.Info("successfully removed provider %s", providerKey)
@@ -3181,6 +3181,15 @@ func (bifrost *Bifrost) RemoveProvider(providerKey schemas.ModelProvider) error
// Note: This operation will temporarily pause request processing for the specified provider
// while the transition occurs. In-flight requests will complete before workers are stopped.
// Buffered requests in the old queue will be transferred to the new queue to prevent loss.
+//
+// Concurrency safety — no-worker window:
+// UpdateProvider holds a per-provider write lock (providerMutex.Lock) for its entire
+// duration. All producer paths (tryRequest, tryStreamRequest) acquire the corresponding
+// read lock inside getProviderQueue before they can look up or enqueue into any queue.
+// This means no producer can observe or enqueue into newPq until UpdateProvider returns
+// and releases the write lock — at which point new workers are already running and
+// consuming newPq. There is therefore no window where newPq is visible to producers
+// but has zero workers.
func (bifrost *Bifrost) UpdateProvider(providerKey schemas.ModelProvider) error {
bifrost.logger.Info(fmt.Sprintf("Updating provider configuration for provider %s", providerKey))
// Get the updated configuration from the account
@@ -3213,23 +3222,23 @@ func (bifrost *Bifrost) UpdateProvider(providerKey schemas.ModelProvider) error
queue: make(chan *ChannelMessage, providerConfig.ConcurrencyAndBufferSize.BufferSize),
done: make(chan struct{}),
signalOnce: sync.Once{},
- closeOnce: sync.Once{},
}
- // Step 2: Atomically replace the queue FIRST (new producers immediately get the new queue)
- // This minimizes the window where requests fail during the update
+ // Step 2: Atomically replace the queue so new producers immediately use newPq.
bifrost.requestQueues.Store(providerKey, newPq)
bifrost.logger.Debug("stored new queue for provider %s, new producers will use it", providerKey)
- // Step 3: Signal old queue is closing to producers that already have a reference
- // Only in-flight producers with the old reference will see this
- oldPq.signalClosing()
- bifrost.logger.Debug("signaled closing for old queue of provider %s", providerKey)
-
- // Step 4: Transfer any buffered requests from old queue to new queue
- // This prevents request loss during the transition
+ // Step 3: Transfer buffered requests from the old queue to the new queue BEFORE
+ // signalling workers to stop. This ensures buffered requests are processed by the
+ // new workers rather than being drained with errors.
+ // Old workers are still running and may consume some items concurrently — that is
+ // fine, they process them normally.
+ // If newPq is full during transfer, all remaining buffered requests are cancelled
+ // immediately rather than blocking — this avoids the deadlock where transfer goroutines
+ // wait for space that only opens once new workers start (which can't happen until
+ // the transfer completes).
transferredCount := 0
- var transferWaitGroup sync.WaitGroup
+ cancelledCount := 0
for {
select {
case msg := <-oldPq.queue:
@@ -3237,37 +3246,33 @@ func (bifrost *Bifrost) UpdateProvider(providerKey schemas.ModelProvider) error
case newPq.queue <- msg:
transferredCount++
default:
- // New queue is full, handle this request in a goroutine
- // This is unlikely with proper buffer sizing but provides safety
- transferWaitGroup.Add(1)
- go func(m *ChannelMessage) {
- defer transferWaitGroup.Done()
+ // newPq is full — cancel this message and all remaining in oldPq.
+ cancelMsg := func(r *ChannelMessage) {
+ prov, mod, _ := r.BifrostRequest.GetRequestFields()
select {
- case newPq.queue <- m:
- // Message successfully transferred
- case <-time.After(5 * time.Second):
- bifrost.logger.Warn("Failed to transfer buffered request to new queue within timeout")
- // Send error response to avoid hanging the client
- provider, model, _ := m.BifrostRequest.GetRequestFields()
- select {
- case m.Err <- schemas.BifrostError{
- IsBifrostError: false,
- Error: &schemas.ErrorField{
- Message: "request failed during provider concurrency update",
- },
- ExtraFields: schemas.BifrostErrorExtraFields{
- RequestType: m.RequestType,
- Provider: provider,
- ModelRequested: model,
- },
- }:
- case <-time.After(1 * time.Second):
- // If we can't send the error either, just log and continue
- bifrost.logger.Warn("Failed to send error response during transfer timeout")
- }
+ case r.Err <- schemas.BifrostError{
+ IsBifrostError: false,
+ Error: &schemas.ErrorField{Message: "request failed during provider concurrency update: queue full"},
+ ExtraFields: schemas.BifrostErrorExtraFields{
+ RequestType: r.RequestType,
+ Provider: prov,
+ ModelRequested: mod,
+ },
+ }:
+ case <-r.Context.Done():
}
- }(msg)
- goto transferComplete
+ }
+ cancelMsg(msg)
+ cancelledCount++
+ for {
+ select {
+ case r := <-oldPq.queue:
+ cancelMsg(r)
+ cancelledCount++
+ default:
+ goto transferComplete
+ }
+ }
}
default:
// No more buffered messages
@@ -3276,33 +3281,59 @@ func (bifrost *Bifrost) UpdateProvider(providerKey schemas.ModelProvider) error
}
transferComplete:
- // Wait for all transfer goroutines to complete
- transferWaitGroup.Wait()
if transferredCount > 0 {
bifrost.logger.Info("transferred %d buffered requests to new queue for provider %s", transferredCount, providerKey)
}
+ if cancelledCount > 0 {
+ bifrost.logger.Warn("cancelled %d buffered requests during transfer for provider %s: new queue was full", cancelledCount, providerKey)
+ }
- // Step 5: Close the old queue to signal workers to stop
- oldPq.closeQueue()
- bifrost.logger.Debug("closed old request queue for provider %s", providerKey)
+ // Step 4: Signal the old queue is closing. Producers that still hold a reference to
+ // oldPq will detect this via isClosing() and transparently re-route to newPq.
+ // This happens after the transfer so the new queue is already populated before
+ // stale producers attempt their re-route.
+ oldPq.signalClosing()
+ bifrost.logger.Debug("signaled closing for old queue of provider %s", providerKey)
- // Step 6: Wait for all existing workers to finish processing in-flight requests
+ // Step 5: Wait for all existing workers to finish processing in-flight requests.
+ // Workers exit via oldPq.done (signalled above).
waitGroup, exists := bifrost.waitGroups.Load(providerKey)
if exists {
waitGroup.(*sync.WaitGroup).Wait()
bifrost.logger.Debug("all workers for provider %s have stopped", providerKey)
}
- // Step 7: Create new wait group for the updated workers
+ // Step 5b: Final drain sweep — see drainQueueWithErrors for full explanation.
+ bifrost.drainQueueWithErrors(oldPq)
+
+ // Step 6: Create new wait group for the updated workers.
bifrost.waitGroups.Store(providerKey, &sync.WaitGroup{})
- // Step 8: Create provider instance
+ // Step 7: Create provider instance.
provider, err := bifrost.createBaseProvider(providerKey, providerConfig)
if err != nil {
- return fmt.Errorf("failed to create provider instance for %s: %v", providerKey, err)
- }
-
- // Step 8.5: Atomically replace the provider in the providers slice
+ // Roll back: signal closing, remove from map, then drain.
+ // Order matters: Delete before drainQueueWithErrors so that producers
+ // re-routing via requestQueues.Load find nothing and return "provider
+ // shutting down" immediately, narrowing the TOCTOU window before the sweep.
+ newPq.signalClosing()
+ bifrost.requestQueues.Delete(providerKey)
+ bifrost.waitGroups.Delete(providerKey)
+ bifrost.drainQueueWithErrors(newPq)
+ if sliceErr := bifrost.removeProviderFromSlice(providerKey); sliceErr != nil {
+ bifrost.logger.Error(
+ "UpdateProvider rollback for %s is incomplete — provider was removed from queues "+
+ "but could not be removed from the providers slice: %v. "+
+ "bifrost.providers is now inconsistent. "+
+ "To recover: call RemoveProvider(%s) then AddProvider to re-register it, "+
+ "or restart Bifrost if that fails.",
+ providerKey, sliceErr, providerKey,
+ )
+ }
+ return fmt.Errorf("provider update for %s failed during initialization; provider has been removed — re-add or retry UpdateProvider to restore it: %v", providerKey, err)
+ }
+
+ // Step 8: Atomically replace the provider in the providers slice.
// This must happen before starting new workers to prevent stale reads
bifrost.logger.Debug("atomically replacing provider instance in providers slice for %s", providerKey)
@@ -3312,7 +3343,21 @@ transferComplete:
for {
replacementAttempts++
if replacementAttempts > maxReplacementAttempts {
- return fmt.Errorf("failed to replace provider %s in providers slice after %d attempts", providerKey, maxReplacementAttempts)
+ newPq.signalClosing()
+ bifrost.requestQueues.Delete(providerKey)
+ bifrost.waitGroups.Delete(providerKey)
+ bifrost.drainQueueWithErrors(newPq)
+ if sliceErr := bifrost.removeProviderFromSlice(providerKey); sliceErr != nil {
+ bifrost.logger.Error(
+ "UpdateProvider rollback for %s is incomplete — provider was removed from queues "+
+ "but could not be removed from the providers slice: %v. "+
+ "bifrost.providers is now inconsistent. "+
+ "To recover: call RemoveProvider(%s) then AddProvider to re-register it, "+
+ "or restart Bifrost if that fails.",
+ providerKey, sliceErr, providerKey,
+ )
+ }
+ return fmt.Errorf("failed to replace provider %s in providers slice after %d attempts; provider has been removed — re-add or retry UpdateProvider to restore it", providerKey, maxReplacementAttempts)
}
oldPtr := bifrost.providers.Load()
@@ -3348,7 +3393,7 @@ transferComplete:
// Retrying as swapping did not work (likely due to concurrent modification)
}
- // Step 9: Start new workers with updated concurrency
+ // Step 9: Start new workers with updated concurrency.
bifrost.logger.Debug("starting %d new workers for provider %s with buffer size %d",
providerConfig.ConcurrencyAndBufferSize.Concurrency,
providerKey,
@@ -3384,6 +3429,33 @@ func (bifrost *Bifrost) getProviderMutex(providerKey schemas.ModelProvider) *syn
return mutexValue.(*sync.RWMutex)
}
+// removeProviderFromSlice atomically removes the provider with the given key
+// from bifrost.providers using a CAS retry loop. Callers hold the per-provider
+// write mutex so no concurrent goroutine can re-add this key — contention is
+// only from other providers' CAS operations, so the loop converges in at most
+// a few iterations under any concurrency level.
+// Returns an error if the limit is hit (state will be inconsistent).
+func (bifrost *Bifrost) removeProviderFromSlice(providerKey schemas.ModelProvider) error {
+ const maxAttempts = 100
+ for range maxAttempts {
+ oldPtr := bifrost.providers.Load()
+ if oldPtr == nil {
+ return nil
+ }
+ oldSlice := *oldPtr
+ newSlice := make([]schemas.Provider, 0, len(oldSlice))
+ for _, p := range oldSlice {
+ if p.GetProviderKey() != providerKey {
+ newSlice = append(newSlice, p)
+ }
+ }
+ if bifrost.providers.CompareAndSwap(oldPtr, &newSlice) {
+ return nil
+ }
+ }
+ return fmt.Errorf("failed to remove provider %s from providers slice after %d attempts", providerKey, maxAttempts)
+}
+
// MCP PUBLIC API
// RegisterMCPTool registers a typed tool handler with the MCP integration.
@@ -3694,7 +3766,6 @@ func (bifrost *Bifrost) prepareProvider(providerKey schemas.ModelProvider, confi
queue: make(chan *ChannelMessage, config.ConcurrencyAndBufferSize.BufferSize),
done: make(chan struct{}),
signalOnce: sync.Once{},
- closeOnce: sync.Once{},
}
bifrost.requestQueues.Store(providerKey, pq)
@@ -4382,17 +4453,31 @@ func (bifrost *Bifrost) tryRequest(ctx *schemas.BifrostContext, req *schemas.Bif
msg := bifrost.getChannelMessage(*preReq)
msg.Context = ctx
- // Check if provider is closing before attempting to send (lock-free atomic check)
- // This prevents "send on closed channel" panics during provider removal/update
+ // If the queue is closing, check whether the provider was updated (new queue
+ // available) or removed. On update, transparently re-route to the new queue
+ // so in-flight producers don't get spurious errors. On removal, error out.
+ //
+ // Use a direct sync.Map lookup instead of getProviderQueue to avoid the
+ // lazy-creation path: getProviderQueue can resurrect a provider that was
+ // just removed by RemoveProvider if the account config still exists.
if pq.isClosing() {
- bifrost.releaseChannelMessage(msg)
- bifrostErr := newBifrostErrorFromMsg("provider is shutting down")
- bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
- RequestType: req.RequestType,
- Provider: provider,
- ModelRequested: model,
+ var reroutedPq *ProviderQueue
+ if val, ok := bifrost.requestQueues.Load(provider); ok {
+ if candidate := val.(*ProviderQueue); candidate != pq && !candidate.isClosing() {
+ reroutedPq = candidate
+ }
}
- return nil, bifrostErr
+ if reroutedPq == nil {
+ bifrost.releaseChannelMessage(msg)
+ bifrostErr := newBifrostErrorFromMsg("provider is shutting down")
+ bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
+ RequestType: req.RequestType,
+ Provider: provider,
+ ModelRequested: model,
+ }
+ return nil, bifrostErr
+ }
+ pq = reroutedPq
}
// Use select with done channel to detect shutdown during send
@@ -4492,7 +4577,13 @@ func (bifrost *Bifrost) tryRequest(ctx *schemas.BifrostContext, req *schemas.Bif
}
return resp, nil
case <-ctx.Done():
- bifrost.releaseChannelMessage(msg)
+ // Do NOT releaseChannelMessage here. The message is already enqueued and
+ // the worker still holds a reference to msg.Response and msg.Err. Returning
+ // those channels to the pool now would let the next request reuse them while
+ // the worker is still writing to them — stale data corruption. The worker
+ // never calls releaseChannelMessage itself, so this message leaks from the
+ // pool and is GC'd. That is intentional: a small pool leak on cancellation
+ // is far safer than corrupting another request's channels.
provider, model, _ := req.GetRequestFields()
return nil, newBifrostCtxDoneError(ctx, provider, model, req.RequestType, "waiting for provider response")
}
@@ -4598,8 +4689,16 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
// shared processedResponse or processedError objects.
streamResponse := providerUtils.BuildClientStreamChunk(ctx, processedResponse, processedError)
- // Send the processed message to the output stream
- outputStream <- streamResponse
+ // Guarded send: if the consumer abandons outputStream (client
+ // disconnect, ctx cancel), drain the upstream shortCircuit.Stream
+ // so its producer can exit cleanly instead of blocking on its send.
+ select {
+ case outputStream <- streamResponse:
+ case <-ctx.Done():
+ for range shortCircuit.Stream {
+ }
+ return
+ }
//TODO: Release the processed response immediately after use
}
@@ -4629,17 +4728,31 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
msg := bifrost.getChannelMessage(*preReq)
msg.Context = ctx
- // Check if provider is closing before attempting to send (lock-free atomic check)
- // This prevents "send on closed channel" panics during provider removal/update
+ // If the queue is closing, check whether the provider was updated (new queue
+ // available) or removed. On update, transparently re-route to the new queue
+ // so in-flight producers don't get spurious errors. On removal, error out.
+ //
+ // Use a direct sync.Map lookup instead of getProviderQueue to avoid the
+ // lazy-creation path: getProviderQueue can resurrect a provider that was
+ // just removed by RemoveProvider if the account config still exists.
if pq.isClosing() {
- bifrost.releaseChannelMessage(msg)
- bifrostErr := newBifrostErrorFromMsg("provider is shutting down")
- bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
- RequestType: req.RequestType,
- Provider: provider,
- ModelRequested: model,
+ var reroutedPq *ProviderQueue
+ if val, ok := bifrost.requestQueues.Load(provider); ok {
+ if candidate := val.(*ProviderQueue); candidate != pq && !candidate.isClosing() {
+ reroutedPq = candidate
+ }
}
- return nil, bifrostErr
+ if reroutedPq == nil {
+ bifrost.releaseChannelMessage(msg)
+ bifrostErr := newBifrostErrorFromMsg("provider is shutting down")
+ bifrostErr.ExtraFields = schemas.BifrostErrorExtraFields{
+ RequestType: req.RequestType,
+ Provider: provider,
+ ModelRequested: model,
+ }
+ return nil, bifrostErr
+ }
+ pq = reroutedPq
}
// Use select with done channel to detect shutdown during send
@@ -4721,6 +4834,11 @@ func (bifrost *Bifrost) tryStreamRequest(ctx *schemas.BifrostContext, req *schem
return newBifrostMessageChan(recoveredResp), nil
}
return nil, &bifrostErrVal
+ case <-ctx.Done():
+ // Do NOT releaseChannelMessage here — see the identical note in tryRequest.
+ // Worker still holds msg.ResponseStream/msg.Err; releasing now corrupts the
+ // next request that reuses those pooled channels.
+ return nil, newBifrostCtxDoneError(ctx, provider, model, req.RequestType, "while waiting for stream response")
}
}
@@ -4842,7 +4960,7 @@ func executeRequestWithRetries[T any](
// the SSE stream instead of returning proper HTTP error status codes.
if bifrostError == nil {
if streamChan, ok := any(result).(chan *schemas.BifrostStreamChunk); ok {
- checkedStream, drainDone, firstChunkErr := providerUtils.CheckFirstStreamChunkForError(streamChan)
+ checkedStream, drainDone, firstChunkErr := providerUtils.CheckFirstStreamChunkForError(ctx, streamChan)
if firstChunkErr != nil {
<-drainDone
bifrostError = firstChunkErr
@@ -4937,7 +5055,38 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
}
}()
- for req := range pq.queue {
+ for {
+ var req *ChannelMessage
+ select {
+ case r := <-pq.queue:
+ req = r
+ case <-pq.done:
+ // Provider is shutting down. Drain any buffered requests and send
+ // back errors so callers are not left blocked on their response channel.
+ for {
+ select {
+ case r := <-pq.queue:
+ provKey, mod, _ := r.GetRequestFields()
+ select {
+ case r.Err <- schemas.BifrostError{
+ IsBifrostError: false,
+ Error: &schemas.ErrorField{
+ Message: "provider is shutting down",
+ },
+ ExtraFields: schemas.BifrostErrorExtraFields{
+ RequestType: r.RequestType,
+ Provider: provKey,
+ ModelRequested: mod,
+ },
+ }:
+ case <-r.Context.Done():
+ }
+ default:
+ return
+ }
+ }
+ }
+
_, model, _ := req.BifrostRequest.GetRequestFields()
var result *schemas.BifrostResponse
@@ -5075,12 +5224,16 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas
}
return resp, nil
}
- // Store a finalizer callback to create aggregated post-hook spans at stream end
- // This closure captures the pipeline reference and releases it after finalization
+ // Store a finalizer callback to create aggregated post-hook spans at stream end.
+ // Wrapped in sync.Once so the normal end-of-stream invocation and a deferred
+ // safety-net invocation (e.g. from a provider goroutine's panic path) cannot
+ // double-release the pipeline.
+ var finalizerOnce sync.Once
postHookSpanFinalizer := func(ctx context.Context) {
- pipeline.FinalizeStreamingPostHookSpans(ctx)
- // Release the pipeline AFTER finalizing spans (not before streaming completes)
- bifrost.releasePluginPipeline(pipeline)
+ finalizerOnce.Do(func() {
+ pipeline.FinalizeStreamingPostHookSpans(ctx)
+ bifrost.releasePluginPipeline(pipeline)
+ })
}
req.Context.SetValue(schemas.BifrostContextKeyPostHookSpanFinalizer, postHookSpanFinalizer)
}
@@ -5206,6 +5359,16 @@ func (bifrost *Bifrost) handleProviderRequest(provider schemas.Provider, config
}
response.RerankResponse = rerankResponse
case schemas.OCRRequest:
+ var customProviderConfig *schemas.CustomProviderConfig
+ if config != nil {
+ customProviderConfig = config.CustomProviderConfig
+ }
+ if bifrostError := providerUtils.CheckOperationAllowed(provider.GetProviderKey(), customProviderConfig, schemas.OCRRequest); bifrostError != nil {
+ if req.BifrostRequest.OCRRequest != nil {
+ bifrostError.ExtraFields.ModelRequested = req.BifrostRequest.OCRRequest.Model
+ }
+ return nil, bifrostError
+ }
ocrResponse, bifrostError := provider.OCR(req.Context, key, req.BifrostRequest.OCRRequest)
if bifrostError != nil {
return nil, bifrostError
@@ -5984,6 +6147,47 @@ func (bifrost *Bifrost) getChannelMessage(req schemas.BifrostRequest) *ChannelMe
return msg
}
+// drainQueueWithErrors drains all buffered messages from pq and sends each a
+// "provider is shutting down" error. It must be called after all workers for
+// the queue have exited (i.e. after wg.Wait()) to cover the TOCTOU window:
+// a producer that passed isClosing() just before signalClosing fired can still
+// win the `case pq.queue <- msg` branch in tryRequest, landing a message in
+// the queue after the last worker's drain loop already exited via `default:`.
+// Without this sweep, those callers block forever on <-msg.Response / <-msg.Err.
+//
+// Residual TOCTOU window (known limitation): this sweep runs exactly once via
+// a non-blocking `select { default: }`. A producer that deposits a message
+// after the sweep's `default:` branch exits has no worker and no sweep to drain
+// it — the caller will block until its own context is cancelled. Fully closing
+// this window requires a sender-side reference count (so the last producer can
+// signal "queue is fully idle"), which is intentionally not implemented because
+// it would add per-send atomic overhead on the hot path.
+func (bifrost *Bifrost) drainQueueWithErrors(pq *ProviderQueue) {
+ for {
+ select {
+ case r := <-pq.queue:
+ provKey, mod, _ := r.GetRequestFields()
+ select {
+ case r.Err <- schemas.BifrostError{
+ IsBifrostError: false,
+ Error: &schemas.ErrorField{Message: "provider is shutting down"},
+ ExtraFields: schemas.BifrostErrorExtraFields{
+ RequestType: r.RequestType,
+ Provider: provKey,
+ ModelRequested: mod,
+ },
+ }:
+ case <-r.Context.Done():
+ // No time.After needed: r.Err is a buffered channel of size 1 freshly
+ // allocated per request, so the send always completes immediately unless
+ // the caller already cancelled. ctx.Done() is the only valid escape.
+ }
+ default:
+ return
+ }
+ }
+}
+
// releaseChannelMessage returns a ChannelMessage and its channels to their respective pools.
func (bifrost *Bifrost) releaseChannelMessage(msg *ChannelMessage) {
// Put channels back in pools
@@ -6491,15 +6695,12 @@ func (bifrost *Bifrost) Shutdown() {
if bifrost.ctx.Err() == nil && bifrost.cancel != nil {
bifrost.cancel()
}
- // ALWAYS close all provider queues to signal workers to stop,
- // even if context was already cancelled. This prevents goroutine leaks.
- // Use the ProviderQueue lifecycle: signal closing, then close the queue
+ // Signal all provider queues to close. Workers exit via pq.done;
+ // we never close pq.queue to avoid "send on closed channel" panics in
+ // producers that are concurrently in tryRequest.
bifrost.requestQueues.Range(func(key, value interface{}) bool {
pq := value.(*ProviderQueue)
- // Signal closing to producers (uses sync.Once internally)
pq.signalClosing()
- // Close the queue to signal workers (uses sync.Once internally)
- pq.closeQueue()
return true
})
@@ -6510,6 +6711,12 @@ func (bifrost *Bifrost) Shutdown() {
return true
})
+ // Final drain sweep — same reasoning as RemoveProvider's Step 3b.
+ bifrost.requestQueues.Range(func(key, value interface{}) bool {
+ bifrost.drainQueueWithErrors(value.(*ProviderQueue))
+ return true
+ })
+
// Cleanup MCP manager
if bifrost.MCPManager != nil {
err := bifrost.MCPManager.Cleanup()
diff --git a/core/bifrost_test.go b/core/bifrost_test.go
index cb22f5e359..6944ed1d9d 100644
--- a/core/bifrost_test.go
+++ b/core/bifrost_test.go
@@ -3,8 +3,10 @@ package bifrost
import (
"context"
"fmt"
+ "runtime"
"strings"
"sync"
+ "sync/atomic"
"testing"
"time"
@@ -1300,3 +1302,998 @@ func TestUpdateProvider_ProviderSliceIntegrity(t *testing.T) {
}
})
}
+
+// TestProviderQueue_SendOnClosedChannel_Race demonstrates the TOCTOU race that
+// caused the "send on closed channel" production panic in the OLD code.
+//
+// The old code called close(pq.queue) during provider shutdown. The sequence:
+// 1. Producer calls isClosing() → false (queue is still open)
+// 2. Concurrently: shutdown calls signalClosing() then close(pq.queue)
+// 3. Producer enters select { case pq.queue <- msg: ... case <-pq.done: ... }
+// → PANIC: Go's selectgo iterates cases in a randomised pollorder. When the
+// closed-channel send case is checked first, it immediately panics via
+// goto sclose — before it can reach the done case.
+// The case <-pq.done: guard only saves you when done happens to be checked
+// first in that random ordering (≈50 % of the time with two cases).
+//
+// THE FIX: pq.queue is never closed. See the ProviderQueue struct comment for
+// the full explanation. This test is kept as a proof-of-concept showing why
+// closing pq.queue is unsafe; the fix is validated by TestProviderQueue_NoPanicWithoutCloseQueue.
+//
+// We run many iterations so that the panic is statistically certain to surface
+// at least once, confirming the hypothesis.
+func TestProviderQueue_SendOnClosedChannel_Race(t *testing.T) {
+ // With two select cases each iteration has a ~50 % chance of panicking.
+ // The probability of never panicking in 200 iterations is (0.5)^200 ≈ 0.
+ const iterations = 200
+ panicCount := 0
+
+ for i := 0; i < iterations; i++ {
+ func() {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ // Synchronization barriers to force the exact race interleaving.
+ passedIsClosingCheck := make(chan struct{})
+ queueClosed := make(chan struct{})
+
+ var panicked bool
+ var wg sync.WaitGroup
+ wg.Add(1)
+
+ // Producer — mirrors the hot path in tryRequest.
+ go func() {
+ defer wg.Done()
+ defer func() {
+ if r := recover(); r != nil && fmt.Sprint(r) == "send on closed channel" {
+ panicked = true
+ }
+ }()
+
+ // Step 1: isClosing() passes — queue is open.
+ if pq.isClosing() {
+ return
+ }
+
+ // Signal: past the isClosing() gate.
+ close(passedIsClosingCheck)
+
+ // Wait for the queue to be closed. This represents the real work
+ // tryRequest does between the isClosing() check and the select
+ // (MCP setup, tracer lookup, plugin pipeline acquisition).
+ <-queueClosed
+
+ // Step 2: enter the exact select guard used in production.
+ // pq.queue is closed AND pq.done is closed.
+ // When selectgo picks the send case first in its random pollorder
+ // it hits goto sclose and panics — the done case cannot save it.
+ msg := &ChannelMessage{}
+ select {
+ case pq.queue <- msg: // panics ~50 % of iterations
+ case <-pq.done: // selected the other ~50 %
+ }
+ }()
+
+ // Closer — mirrors UpdateProvider / RemoveProvider.
+ go func() {
+ <-passedIsClosingCheck
+ pq.signalClosing() // closes done, sets closing = 1
+ close(pq.queue)
+ close(queueClosed) // release the producer into the select
+ }()
+
+ wg.Wait()
+ if panicked {
+ panicCount++
+ }
+ }()
+ }
+
+ if panicCount == 0 {
+ t.Fatalf("expected at least one 'send on closed channel' panic across %d iterations, got none", iterations)
+ }
+ t.Logf("confirmed: panic triggered in %d / %d iterations — hypothesis is correct", panicCount, iterations)
+}
+
+// =============================================================================
+// ProviderQueue Unit Tests
+//
+// These tests exercise the ProviderQueue lifecycle in isolation — no full
+// Bifrost instance required. They validate the core safety invariants that
+// prevent the "send on closed channel" panic.
+// =============================================================================
+
+// newTestChannelMessage creates a minimal ChannelMessage suitable for drain tests.
+// The Err channel is buffered (size 1) so the worker can send without blocking.
+func newTestChannelMessage(ctx *schemas.BifrostContext) *ChannelMessage {
+ return &ChannelMessage{
+ BifrostRequest: schemas.BifrostRequest{
+ RequestType: schemas.ChatCompletionRequest,
+ ChatRequest: &schemas.BifrostChatRequest{
+ Provider: schemas.OpenAI,
+ Model: "gpt-4",
+ },
+ },
+ Context: ctx,
+ Response: make(chan *schemas.BifrostResponse, 1),
+ Err: make(chan schemas.BifrostError, 1),
+ }
+}
+
+// TestProviderQueue_IsClosingStateTransition verifies the atomic state flag:
+// isClosing() must return false before signalClosing() and true after.
+func TestProviderQueue_IsClosingStateTransition(t *testing.T) {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ if pq.isClosing() {
+ t.Fatal("isClosing() must be false before signalClosing() is called")
+ }
+
+ pq.signalClosing()
+
+ if !pq.isClosing() {
+ t.Fatal("isClosing() must be true after signalClosing() is called")
+ }
+
+ // done channel must also be closed
+ select {
+ case <-pq.done:
+ // correct: done is closed
+ default:
+ t.Fatal("pq.done must be closed after signalClosing()")
+ }
+
+ // queue channel must remain OPEN — this is the core of the fix
+ // (sending should not panic even though done is closed)
+ panicked := false
+ func() {
+ defer func() {
+ if r := recover(); r != nil {
+ panicked = true
+ }
+ }()
+ select {
+ case pq.queue <- &ChannelMessage{}:
+ case <-pq.done: // done is closed so this is always ready — no panic
+ }
+ }()
+ if panicked {
+ t.Fatal("queue channel must stay open after signalClosing() — sending to it must not panic")
+ }
+}
+
+// TestProviderQueue_SignalOnceIdempotent verifies that calling signalClosing()
+// multiple times is safe. sync.Once ensures done is only closed once and the
+// atomic store only happens once — no "close of closed channel" panic.
+func TestProviderQueue_SignalOnceIdempotent(t *testing.T) {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ defer func() {
+ if r := recover(); r != nil {
+ t.Fatalf("unexpected panic from multiple signalClosing() calls: %v", r)
+ }
+ }()
+
+ pq.signalClosing()
+ pq.signalClosing()
+ pq.signalClosing()
+
+ if !pq.isClosing() {
+ t.Fatal("isClosing() must be true after multiple signalClosing() calls")
+ }
+}
+
+// TestProviderQueue_WorkerExitsViaDone verifies that a worker running the
+// fixed select loop exits cleanly after signalClosing() without closeQueue().
+// Before the fix, workers used `for req := range pq.queue` which required
+// the channel to be closed. After the fix, done is the exit signal.
+func TestProviderQueue_WorkerExitsViaDone(t *testing.T) {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ workerExited := make(chan struct{})
+
+ // Minimal worker loop — mirrors the exact select pattern in requestWorker
+ go func() {
+ defer close(workerExited)
+ for {
+ select {
+ case r, ok := <-pq.queue:
+ if !ok {
+ return
+ }
+ _ = r // process (no-op in this test)
+ case <-pq.done:
+ // Drain remaining buffered items (queue is empty here)
+ for {
+ select {
+ case <-pq.queue:
+ default:
+ return
+ }
+ }
+ }
+ }
+ }()
+
+ // Worker is now blocked on the select. Signal shutdown WITHOUT closing queue.
+ pq.signalClosing()
+
+ select {
+ case <-workerExited:
+ // correct: worker exited via done
+ case <-time.After(2 * time.Second):
+ t.Fatal("worker did not exit after signalClosing() — it may be stuck on range over unclosed channel")
+ }
+}
+
+// TestProviderQueue_WorkerDrainSendsErrors verifies the drain behaviour when
+// done fires while items are still buffered: every buffered ChannelMessage must
+// receive a "provider is shutting down" error on its Err channel. No client
+// should be left blocked waiting for a response that will never come.
+//
+// This test exercises the drain path directly — same code as requestWorker's
+// case <-pq.done: branch — to avoid a non-deterministic select race between the
+// normal processing path and the done path.
+func TestProviderQueue_WorkerDrainSendsErrors(t *testing.T) {
+ const numBuffered = 5
+
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, numBuffered+2),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+
+ // Pre-fill queue — simulates requests buffered when done fires
+ msgs := make([]*ChannelMessage, numBuffered)
+ for i := 0; i < numBuffered; i++ {
+ msgs[i] = newTestChannelMessage(ctx)
+ pq.queue <- msgs[i]
+ }
+
+ // Signal closing: done is now closed
+ pq.signalClosing()
+
+ // Execute the drain path synchronously — exactly what requestWorker does in
+ // the case <-pq.done: branch. This is deterministic: we know done is closed
+ // and the queue has numBuffered items.
+ <-pq.done // fires immediately since signalClosing was already called
+drainLoop:
+ for {
+ select {
+ case r := <-pq.queue:
+ provKey, mod, _ := r.GetRequestFields()
+ r.Err <- schemas.BifrostError{
+ IsBifrostError: false,
+ Error: &schemas.ErrorField{
+ Message: "provider is shutting down",
+ },
+ ExtraFields: schemas.BifrostErrorExtraFields{
+ RequestType: r.RequestType,
+ Provider: provKey,
+ ModelRequested: mod,
+ },
+ }
+ default:
+ break drainLoop
+ }
+ }
+
+ // Verify every message received a shutdown error
+ for i, msg := range msgs {
+ select {
+ case bifrostErr := <-msg.Err:
+ if bifrostErr.Error == nil {
+ t.Errorf("message %d: received nil Error field", i)
+ continue
+ }
+ if bifrostErr.Error.Message != "provider is shutting down" {
+ t.Errorf("message %d: expected 'provider is shutting down', got %q",
+ i, bifrostErr.Error.Message)
+ }
+ if bifrostErr.ExtraFields.Provider != schemas.OpenAI {
+ t.Errorf("message %d: expected provider %s, got %s",
+ i, schemas.OpenAI, bifrostErr.ExtraFields.Provider)
+ }
+ if bifrostErr.ExtraFields.RequestType != schemas.ChatCompletionRequest {
+ t.Errorf("message %d: expected requestType %v, got %v",
+ i, schemas.ChatCompletionRequest, bifrostErr.ExtraFields.RequestType)
+ }
+ default:
+ t.Errorf("message %d: no error received — client would be left hanging indefinitely", i)
+ }
+ }
+}
+
+// TestProviderQueue_NoPanicWithoutCloseQueue verifies that the fixed hot path
+// — select { case pq.queue <- msg | case <-pq.done } — never panics when
+// signalClosing() fires but the queue channel is NOT closed.
+//
+// This is the direct inverse of TestProviderQueue_SendOnClosedChannel_Race:
+// that test proves the old code panics ~50% of the time; this test proves
+// the fixed code panics 0% of the time.
+func TestProviderQueue_NoPanicWithoutCloseQueue(t *testing.T) {
+ const iterations = 500
+
+ for i := 0; i < iterations; i++ {
+ func() {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ passedIsClosingCheck := make(chan struct{})
+ shutdownDone := make(chan struct{})
+
+ var panicked bool
+ var wg sync.WaitGroup
+ wg.Add(1)
+
+ // Producer: mirrors the tryRequest hot path after the fix.
+ // Passes isClosing(), waits for signalClosing, then sends.
+ // The queue channel is NEVER closed — only done is closed.
+ go func() {
+ defer wg.Done()
+ defer func() {
+ if r := recover(); r != nil {
+ panicked = true
+ }
+ }()
+
+ if pq.isClosing() {
+ return
+ }
+ close(passedIsClosingCheck)
+ <-shutdownDone
+
+ msg := &ChannelMessage{}
+ select {
+ case pq.queue <- msg: // queue is open → safe to send
+ case <-pq.done: // done is closed → selected immediately
+ }
+ }()
+
+ // Closer: signal shutdown but never close the queue channel
+ go func() {
+ <-passedIsClosingCheck
+ pq.signalClosing() // closes done; does NOT close queue
+ close(shutdownDone)
+ }()
+
+ wg.Wait()
+
+ if panicked {
+ t.Errorf("iteration %d: unexpected panic — queue must not be closed in the fixed path", i)
+ }
+ }()
+
+ if t.Failed() {
+ return
+ }
+ }
+
+ t.Logf("confirmed: zero panics in %d iterations with the fix applied", iterations)
+}
+
+// =============================================================================
+// UpdateProvider Lifecycle Tests
+//
+// These tests verify the three key invariants of the UpdateProvider fix:
+// 1. New queue is stored BEFORE signalClosing fires (stale producers re-route)
+// 2. Transfer happens BEFORE signalClosing (items go to new workers, not errored)
+// 3. Concurrent producers + UpdateProvider produce zero panics
+// =============================================================================
+
+// TestUpdateProvider_StaleProducerReroutes verifies that a "stale producer" —
+// a goroutine that fetched oldPq before UpdateProvider atomically replaced it —
+// can transparently re-route to newPq when it later detects isClosing().
+//
+// The re-routing logic in tryRequest is:
+//
+// if pq.isClosing() {
+// if newPq, err := bifrost.getProviderQueue(provider); err == nil && newPq != pq {
+// pq = newPq // transparent re-route
+// }
+// }
+//
+// This test exercises that exact sequence without a full Bifrost instance.
+func TestUpdateProvider_StaleProducerReroutes(t *testing.T) {
+ var requestQueues sync.Map
+ provider := schemas.OpenAI
+
+ oldPq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+ newPq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ // Initial state: requestQueues holds oldPq
+ requestQueues.Store(provider, oldPq)
+
+ // Stale producer: fetched its reference before UpdateProvider ran
+ stalePq := oldPq
+
+ // Simulate UpdateProvider steps 2 + 4:
+ // Step 2: atomically replace — new producers now get newPq
+ requestQueues.Store(provider, newPq)
+ // Step 4: signal old closing — stale producers will detect this
+ oldPq.signalClosing()
+
+ // --- Stale producer detects isClosing and attempts re-route ---
+ var reroutedPq *ProviderQueue
+ if stalePq.isClosing() {
+ if val, ok := requestQueues.Load(provider); ok {
+ candidate := val.(*ProviderQueue)
+ if candidate != stalePq {
+ reroutedPq = candidate
+ }
+ }
+ }
+
+ if reroutedPq == nil {
+ t.Fatal("stale producer failed to re-route: re-route returned nil (check step ordering)")
+ }
+ if reroutedPq != newPq {
+ t.Fatal("stale producer re-routed to wrong queue: expected newPq")
+ }
+ if reroutedPq.isClosing() {
+ t.Fatal("re-routed queue is already closing — re-route is useless (newPq must be fresh)")
+ }
+
+ // Verify: sending to re-routed queue succeeds without panic
+ panicked := false
+ func() {
+ defer func() {
+ if r := recover(); r != nil {
+ panicked = true
+ }
+ }()
+ msg := &ChannelMessage{}
+ select {
+ case reroutedPq.queue <- msg:
+ case <-reroutedPq.done:
+ t.Error("newPq.done fired — newPq should be open")
+ }
+ }()
+ if panicked {
+ t.Fatal("panic while sending to re-routed queue — queue must not be closed")
+ }
+}
+
+// TestUpdateProvider_TransferOrdering verifies the ordering invariant:
+// items are moved from oldPq to newPq BEFORE signalClosing(oldPq) is called.
+//
+// Observable consequence: during the entire transfer loop, oldPq.isClosing()
+// must remain false. Only after transfer completes does signalClosing fire.
+func TestUpdateProvider_TransferOrdering(t *testing.T) {
+ const numMessages = 8
+
+ oldPq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, numMessages+2),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+ newPq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, numMessages+2),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ // Pre-fill oldPq — simulates buffered requests at the moment UpdateProvider runs
+ for i := 0; i < numMessages; i++ {
+ oldPq.queue <- &ChannelMessage{}
+ }
+
+ // Invariant check before transfer begins
+ if oldPq.isClosing() {
+ t.Fatal("invariant violated: oldPq already closing before transfer begins")
+ }
+
+ // Perform transfer, mirroring UpdateProvider step 3.
+ // Record whether isClosing() ever fired during the loop.
+ closingDuringTransfer := false
+ transferred := 0
+ for {
+ select {
+ case msg := <-oldPq.queue:
+ if oldPq.isClosing() {
+ closingDuringTransfer = true
+ }
+ newPq.queue <- msg
+ transferred++
+ default:
+ goto transferComplete
+ }
+ }
+transferComplete:
+
+ if closingDuringTransfer {
+ t.Error("invariant violated: oldPq was already closing during transfer — " +
+ "signalClosing must fire AFTER the transfer loop completes")
+ }
+
+ // NOW signal closing, mirroring UpdateProvider step 4
+ oldPq.signalClosing()
+
+ if !oldPq.isClosing() {
+ t.Error("expected isClosing() == true after signalClosing()")
+ }
+
+ // All messages must have moved to newPq
+ if transferred != numMessages {
+ t.Errorf("expected %d messages transferred, got %d", numMessages, transferred)
+ }
+ if len(newPq.queue) != numMessages {
+ t.Errorf("expected %d messages in newPq after transfer, got %d", numMessages, len(newPq.queue))
+ }
+ if len(oldPq.queue) != 0 {
+ t.Errorf("expected 0 messages remaining in oldPq after transfer, got %d", len(oldPq.queue))
+ }
+}
+
+// TestUpdateProvider_NoPanicConcurrentAccess verifies that concurrent producers
+// sending to a queue that is being replaced (UpdateProvider-style) never cause
+// a "send on closed channel" panic.
+//
+// This test directly models the production scenario that triggered the bug:
+// many goroutines continuously send to a ProviderQueue while UpdateProvider
+// atomically swaps the queue and signals the old one closing. With the fix
+// (queue channel is never closed), the select in producers is always safe.
+func TestUpdateProvider_NoPanicConcurrentAccess(t *testing.T) {
+ const (
+ numProducers = 10
+ numUpdates = 30
+ producerRunTime = 300 * time.Millisecond
+ )
+
+ var requestQueues sync.Map
+ provider := schemas.OpenAI
+
+ makePq := func() *ProviderQueue {
+ return &ProviderQueue{
+ queue: make(chan *ChannelMessage, 200),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+ }
+
+ initialPq := makePq()
+ requestQueues.Store(provider, initialPq)
+
+ var panicCount int64
+ var transferDropCount int64
+
+ stop := make(chan struct{})
+ var producerWg sync.WaitGroup
+
+ // Drainer: continuously empties queues so producers never block on a full queue
+ drainStop := make(chan struct{})
+ go func() {
+ for {
+ select {
+ case <-drainStop:
+ return
+ default:
+ if val, ok := requestQueues.Load(provider); ok {
+ pq := val.(*ProviderQueue)
+ select {
+ case <-pq.queue:
+ default:
+ }
+ }
+ runtime.Gosched()
+ }
+ }
+ }()
+
+ // Producers: continuously simulate the tryRequest hot path
+ for i := 0; i < numProducers; i++ {
+ producerWg.Add(1)
+ go func() {
+ defer producerWg.Done()
+ for {
+ select {
+ case <-stop:
+ return
+ default:
+ }
+
+ val, ok := requestQueues.Load(provider)
+ if !ok {
+ runtime.Gosched()
+ continue
+ }
+ pq := val.(*ProviderQueue)
+
+ func() {
+ defer func() {
+ if r := recover(); r != nil {
+ atomic.AddInt64(&panicCount, 1)
+ }
+ }()
+
+ // Re-route check (mirrors tryRequest)
+ if pq.isClosing() {
+ if newVal, ok2 := requestQueues.Load(provider); ok2 {
+ if candidate := newVal.(*ProviderQueue); candidate != pq {
+ pq = candidate
+ }
+ }
+ // If still closing (RemoveProvider path), just return
+ if pq.isClosing() {
+ return
+ }
+ }
+
+ msg := &ChannelMessage{}
+ select {
+ case pq.queue <- msg:
+ case <-pq.done:
+ case <-stop: // unblock immediately when the test signals stop
+ }
+ }()
+
+ runtime.Gosched()
+ }
+ }()
+ }
+
+ // Updater: repeatedly performs UpdateProvider-style queue replacements
+ var updaterWg sync.WaitGroup
+ updaterWg.Add(1)
+ go func() {
+ defer updaterWg.Done()
+ for i := 0; i < numUpdates; i++ {
+ val, ok := requestQueues.Load(provider)
+ if !ok {
+ continue
+ }
+ oldPq := val.(*ProviderQueue)
+ newPq := makePq()
+
+ // Mirror production UpdateProvider step order exactly:
+ // Step 2: expose newPq first so stale producers can re-route to it
+ // once they see oldPq is closing.
+ requestQueues.Store(provider, newPq)
+
+ // Step 3: transfer buffered messages oldPq → newPq.
+ drain:
+ for {
+ select {
+ case msg := <-oldPq.queue:
+ select {
+ case newPq.queue <- msg:
+ default:
+ // newPq full during transfer — mirrors production cancel path.
+ atomic.AddInt64(&transferDropCount, 1)
+ }
+ default:
+ break drain
+ }
+ }
+
+ // Step 4: signal closing — producers holding a stale oldPq ref now
+ // re-route to newPq (already in the map from step 2).
+ oldPq.signalClosing()
+
+ time.Sleep(5 * time.Millisecond)
+ }
+ }()
+
+ time.Sleep(producerRunTime)
+ close(stop)
+ close(drainStop)
+ producerWg.Wait()
+ updaterWg.Wait()
+
+ if n := atomic.LoadInt64(&panicCount); n > 0 {
+ t.Errorf("detected %d panic(s) — fix did not eliminate the concurrent-access race", n)
+ } else {
+ t.Logf("confirmed: zero panics across %d producers + %d queue replacements over %v",
+ numProducers, numUpdates, producerRunTime)
+ }
+ if drops := atomic.LoadInt64(&transferDropCount); drops > 0 {
+ t.Logf("note: %d message(s) dropped during transfer (oldPq had >200 buffered items) — does not affect panic correctness", drops)
+ }
+}
+
+// =============================================================================
+// RemoveProvider Lifecycle Tests
+//
+// These tests verify the behavioral contract of RemoveProvider:
+// 1. signalClosing() blocks new producers (isClosing() → true)
+// 2. Buffered items in the queue get "provider is shutting down" errors
+// 3. Workers exit cleanly and the WaitGroup reaches zero
+// =============================================================================
+
+// TestRemoveProvider_BlocksNewProducers verifies that after signalClosing(),
+// isClosing() returns true. Producers check this flag before sending and return
+// a "provider is shutting down" error rather than trying to enqueue.
+func TestRemoveProvider_BlocksNewProducers(t *testing.T) {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ // Sanity: before shutdown, producers can proceed
+ if pq.isClosing() {
+ t.Fatal("isClosing() must be false before RemoveProvider runs")
+ }
+
+ // RemoveProvider step 2: signal closing
+ pq.signalClosing()
+
+ // New producers must see isClosing() == true and abort
+ if !pq.isClosing() {
+ t.Fatal("isClosing() must be true after signalClosing() (RemoveProvider)")
+ }
+
+ // done must be closed so any producer blocked in the select unblocks immediately
+ select {
+ case <-pq.done:
+ // correct
+ default:
+ t.Fatal("pq.done must be closed after signalClosing() so blocking producers unblock")
+ }
+
+ // CRITICAL: queue channel must remain OPEN — closing it would cause panics in
+ // any producer that entered the select before seeing isClosing().
+ // With the fix, we NEVER close the queue channel.
+ panicked := false
+ func() {
+ defer func() {
+ if r := recover(); r != nil {
+ panicked = true
+ }
+ }()
+ // A select with done closed always takes the done case — safe, no panic
+ select {
+ case pq.queue <- &ChannelMessage{}:
+ case <-pq.done:
+ }
+ }()
+ if panicked {
+ t.Fatal("queue channel must stay open after signalClosing() — closing it causes panics")
+ }
+}
+
+// TestRemoveProvider_BufferedRequestsGetErrors verifies the drain contract:
+// items queued BEFORE signalClosing fires must each receive a
+// "provider is shutting down" error on their Err channel. No client should be
+// left hanging.
+//
+// This test exercises the drain logic directly — the same code path that
+// requestWorker executes in its case <-pq.done: branch — to avoid the
+// non-deterministic select race where the normal processing path can pick up
+// items before done fires.
+func TestRemoveProvider_BufferedRequestsGetErrors(t *testing.T) {
+ const numBuffered = 8
+
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, numBuffered+5),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline)
+
+ // Buffer requests — simulates requests already queued when RemoveProvider runs
+ msgs := make([]*ChannelMessage, numBuffered)
+ for i := 0; i < numBuffered; i++ {
+ msgs[i] = newTestChannelMessage(ctx)
+ pq.queue <- msgs[i]
+ }
+
+ // RemoveProvider step 2: signal closing
+ pq.signalClosing()
+
+ // Execute the drain path — exactly what requestWorker does in case <-pq.done:
+ <-pq.done // fires immediately since signalClosing was already called
+drainLoop:
+ for {
+ select {
+ case r := <-pq.queue:
+ provKey, mod, _ := r.GetRequestFields()
+ r.Err <- schemas.BifrostError{
+ IsBifrostError: false,
+ Error: &schemas.ErrorField{
+ Message: "provider is shutting down",
+ },
+ ExtraFields: schemas.BifrostErrorExtraFields{
+ RequestType: r.RequestType,
+ Provider: provKey,
+ ModelRequested: mod,
+ },
+ }
+ default:
+ break drainLoop
+ }
+ }
+
+ // Every buffered message must have received a shutdown error
+ for i, msg := range msgs {
+ select {
+ case bifrostErr := <-msg.Err:
+ if bifrostErr.Error == nil {
+ t.Errorf("message %d: got nil Error field in BifrostError", i)
+ continue
+ }
+ if bifrostErr.Error.Message != "provider is shutting down" {
+ t.Errorf("message %d: expected 'provider is shutting down', got %q",
+ i, bifrostErr.Error.Message)
+ }
+ if bifrostErr.ExtraFields.Provider != schemas.OpenAI {
+ t.Errorf("message %d: expected provider %s, got %s",
+ i, schemas.OpenAI, bifrostErr.ExtraFields.Provider)
+ }
+ if bifrostErr.ExtraFields.RequestType != schemas.ChatCompletionRequest {
+ t.Errorf("message %d: expected requestType %v, got %v",
+ i, schemas.ChatCompletionRequest, bifrostErr.ExtraFields.RequestType)
+ }
+ default:
+ t.Errorf("message %d: no error received — client would be left hanging indefinitely", i)
+ }
+ }
+}
+
+// TestRemoveProvider_WorkerWaitGroupCompletes verifies that after signalClosing(),
+// the worker goroutine decrements the WaitGroup and wg.Wait() returns promptly.
+// This mirrors what RemoveProvider does: signal, then Wait() before cleanup.
+func TestRemoveProvider_WorkerWaitGroupCompletes(t *testing.T) {
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, 10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ var wg sync.WaitGroup
+ wg.Add(1)
+
+ // Worker goroutine — mirrors requestWorker's WaitGroup contract
+ go func() {
+ defer wg.Done()
+ for {
+ select {
+ case r, ok := <-pq.queue:
+ if !ok {
+ return
+ }
+ _ = r
+ case <-pq.done:
+ // Drain remaining (empty in this test)
+ for {
+ select {
+ case <-pq.queue:
+ default:
+ return
+ }
+ }
+ }
+ }
+ }()
+
+ // Tiny sleep to ensure worker is parked on select before we signal
+ time.Sleep(10 * time.Millisecond)
+
+ // RemoveProvider step 2: signal closing
+ pq.signalClosing()
+
+ // RemoveProvider step 3: wait for workers — must complete promptly
+ waitReturned := make(chan struct{})
+ go func() {
+ wg.Wait()
+ close(waitReturned)
+ }()
+
+ select {
+ case <-waitReturned:
+ // correct: WaitGroup reached zero after signalClosing()
+ case <-time.After(2 * time.Second):
+ t.Fatal("wg.Wait() did not return after signalClosing() — worker is stuck (would deadlock RemoveProvider)")
+ }
+}
+
+// TestRemoveProvider_ConcurrentNewProducersDuringShutdown verifies that
+// concurrent producers trying to enqueue after RemoveProvider calls
+// signalClosing() all get safe "provider is shutting down" errors — none panic.
+// This tests the TOCTOU window: producer passes isClosing() check, then done fires.
+func TestRemoveProvider_ConcurrentNewProducersDuringShutdown(t *testing.T) {
+ const numProducers = 50
+
+ pq := &ProviderQueue{
+ queue: make(chan *ChannelMessage, numProducers+10),
+ done: make(chan struct{}),
+ signalOnce: sync.Once{},
+ }
+
+ var panicCount int64
+ var shutdownErrors int64
+ var successfulSends int64
+
+ // Gate: all producers start together after isClosing() passes
+ passedGate := make(chan struct{})
+ var gateOnce sync.Once
+ shutdownFired := make(chan struct{})
+
+ var producerWg sync.WaitGroup
+
+ for i := 0; i < numProducers; i++ {
+ producerWg.Add(1)
+ go func() {
+ defer producerWg.Done()
+ defer func() {
+ if r := recover(); r != nil {
+ atomic.AddInt64(&panicCount, 1)
+ }
+ }()
+
+ // Each producer checks isClosing() first (mirrors tryRequest)
+ if pq.isClosing() {
+ atomic.AddInt64(&shutdownErrors, 1)
+ return
+ }
+
+ // Signal that at least one producer passed the isClosing() check
+ gateOnce.Do(func() { close(passedGate) })
+
+ // Wait for shutdown to be signaled (the TOCTOU window)
+ <-shutdownFired
+
+ // Producers now enter the select — with the fix, done is closed but
+ // queue is NOT closed, so this select is always safe (no panic)
+ msg := &ChannelMessage{}
+ select {
+ case pq.queue <- msg:
+ atomic.AddInt64(&successfulSends, 1)
+ case <-pq.done:
+ atomic.AddInt64(&shutdownErrors, 1)
+ }
+ }()
+ }
+
+ // Wait for at least one producer to pass the isClosing() gate
+ select {
+ case <-passedGate:
+ case <-time.After(2 * time.Second):
+ t.Fatal("no producer passed the isClosing() check within timeout")
+ }
+
+ // Signal shutdown (RemoveProvider step 2) — this is the TOCTOU race
+ pq.signalClosing()
+ close(shutdownFired)
+
+ producerWg.Wait()
+
+ if n := atomic.LoadInt64(&panicCount); n > 0 {
+ t.Errorf("detected %d panic(s) — queue must not be closed during concurrent shutdown", n)
+ }
+
+ t.Logf("result: %d successful sends, %d shutdown errors, %d panics across %d producers",
+ atomic.LoadInt64(&successfulSends),
+ atomic.LoadInt64(&shutdownErrors),
+ atomic.LoadInt64(&panicCount),
+ numProducers)
+}
diff --git a/core/go.mod b/core/go.mod
index 924e204ea2..b85c403ec6 100644
--- a/core/go.mod
+++ b/core/go.mod
@@ -1,6 +1,6 @@
module github.com/maximhq/bifrost/core
-go 1.26.2
+go 1.26.1
require (
cloud.google.com/go v0.123.0
diff --git a/core/internal/llmtests/account.go b/core/internal/llmtests/account.go
index ac850830fd..0632f16b6b 100644
--- a/core/internal/llmtests/account.go
+++ b/core/internal/llmtests/account.go
@@ -88,7 +88,9 @@ type TestScenarios struct {
Realtime bool // Realtime API (bidirectional audio/text)
Compaction bool // Server-side compaction (context management)
InterleavedThinking bool // Interleaved thinking between tool calls (beta)
- FastMode bool // Fast mode for Opus 4.6 (beta: research preview)
+ FastMode bool // Fast mode for Opus 4.6 (beta: research preview)
+ EagerInputStreaming bool // Fine-grained tool input streaming (Anthropic fine-grained-tool-streaming-2025-05-14)
+ ServerToolsViaOpenAIEndpoint bool // Anthropic server-tool shapes in tools[] via /v1/chat/completions (web_search / web_fetch / code_execution)
}
// ComprehensiveTestConfig extends TestConfig with additional scenarios
diff --git a/core/internal/llmtests/eager_input_streaming.go b/core/internal/llmtests/eager_input_streaming.go
new file mode 100644
index 0000000000..0f074c46af
--- /dev/null
+++ b/core/internal/llmtests/eager_input_streaming.go
@@ -0,0 +1,134 @@
+package llmtests
+
+import (
+ "context"
+ "os"
+ "testing"
+
+ bifrost "github.com/maximhq/bifrost/core"
+ "github.com/maximhq/bifrost/core/schemas"
+)
+
+// RunEagerInputStreamingTest tests that setting eager_input_streaming: true on
+// a custom tool succeeds end-to-end against the target Anthropic-family
+// provider. Per Table 20 (verified against A overview + B-header), the
+// fine-grained-tool-streaming-2025-05-14 beta is supported on Anthropic,
+// Bedrock, Vertex, and Azure.
+//
+// The test verifies:
+// 1. The request is accepted (no upstream 400 — which would indicate the
+// fine-grained-tool-streaming-2025-05-14 beta header wasn't injected or
+// is rejected by the target provider).
+// 2. The stream produces a tool call with a valid JSON arguments payload.
+// 3. The response is otherwise well-formed.
+//
+// This intentionally runs across all four providers (no single-provider gate
+// unlike RunFastModeTest, which is Opus-4.6-only).
+func RunEagerInputStreamingTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+ if !testConfig.Scenarios.EagerInputStreaming {
+ t.Logf("EagerInputStreaming not supported for provider %s", testConfig.Provider)
+ return
+ }
+
+ t.Run("EagerInputStreaming", func(t *testing.T) {
+ if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
+ t.Parallel()
+ }
+
+ chatTool := GetSampleChatTool(SampleToolTypeWeather)
+ // Opt the tool into fine-grained input streaming. The neutral flag
+ // on ChatTool is promoted through ToAnthropicChatRequest, which also
+ // triggers the fine-grained-tool-streaming-2025-05-14 beta header.
+ eager := true
+ chatTool.EagerInputStreaming = &eager
+
+ chatMessages := []schemas.ChatMessage{
+ CreateBasicChatMessage("What's the weather like in San Francisco? answer in celsius"),
+ }
+
+ request := &schemas.BifrostChatRequest{
+ Provider: testConfig.Provider,
+ Model: testConfig.ChatModel,
+ Input: chatMessages,
+ Params: &schemas.ChatParameters{
+ MaxCompletionTokens: bifrost.Ptr(200),
+ Tools: []schemas.ChatTool{*chatTool},
+ },
+ Fallbacks: testConfig.Fallbacks,
+ }
+
+ retryConfig := StreamingRetryConfig()
+ retryContext := TestRetryContext{
+ ScenarioName: "EagerInputStreaming",
+ ExpectedBehavior: map[string]interface{}{
+ "should_stream_content": true,
+ "should_have_tool_calls": true,
+ "tool_name": "get_weather",
+ },
+ TestMetadata: map[string]interface{}{
+ "provider": testConfig.Provider,
+ "model": testConfig.ChatModel,
+ "eager_input_streaming": true,
+ },
+ }
+
+ responseChannel, err := WithStreamRetry(t, retryConfig, retryContext, func() (chan *schemas.BifrostStreamChunk, *schemas.BifrostError) {
+ bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline)
+ return client.ChatCompletionStreamRequest(bfCtx, request)
+ })
+
+ RequireNoError(t, err, "Eager input streaming request failed")
+ if responseChannel == nil {
+ t.Fatal("Response channel should not be nil")
+ }
+
+ accumulator := NewStreamingToolCallAccumulator()
+ var responseCount int
+ var sawAny bool
+
+ t.Logf("🔧 Testing eager input streaming (fine-grained-tool-streaming-2025-05-14)...")
+
+ for response := range responseChannel {
+ if response == nil || response.BifrostChatResponse == nil {
+ continue
+ }
+ responseCount++
+ sawAny = true
+
+ if response.BifrostChatResponse.Choices != nil {
+ for i, choice := range response.BifrostChatResponse.Choices {
+ if choice.ChatStreamResponseChoice != nil && choice.ChatStreamResponseChoice.Delta != nil {
+ delta := choice.ChatStreamResponseChoice.Delta
+ for _, tc := range delta.ToolCalls {
+ accumulator.AccumulateChatToolCall(i, tc)
+ }
+ }
+ }
+ }
+ }
+
+ if !sawAny {
+ t.Fatal("Expected at least one streaming response chunk")
+ }
+ t.Logf("Received %d chunks", responseCount)
+
+ // Validate the accumulated tool call is well-formed. If the
+ // fine-grained-tool-streaming beta header weren't sent (or the
+ // provider rejected it), the upstream would have returned a 400
+ // before any tool_use blocks were emitted.
+ toolCalls := accumulator.GetFinalChatToolCalls()
+ if len(toolCalls) == 0 {
+ t.Error("Expected at least one tool call in stream")
+ }
+ for _, tc := range toolCalls {
+ if tc.Name == "" {
+ t.Error("Tool call missing function name")
+ }
+ if tc.Arguments == "" {
+ t.Error("Tool call missing arguments JSON")
+ }
+ }
+
+ t.Logf("EagerInputStreaming passed: %d tool calls accumulated", len(toolCalls))
+ })
+}
diff --git a/core/internal/llmtests/provider_feature_support_test.go b/core/internal/llmtests/provider_feature_support_test.go
index 6e4738c282..539f4049c0 100644
--- a/core/internal/llmtests/provider_feature_support_test.go
+++ b/core/internal/llmtests/provider_feature_support_test.go
@@ -654,6 +654,77 @@ func TestProviderBetaHeaderInjection(t *testing.T) {
},
expectHeaders: []string{"computer-use-2025-01-24"},
},
+
+ // ── Fine-grained tool streaming header (eager_input_streaming) ──
+ // Per cited citations (A overview table + B-header): EagerInputStreaming
+ // is supported on Anthropic, Bedrock, Vertex, and Azure — all four
+ // should auto-inject fine-grained-tool-streaming-2025-05-14 when a
+ // tool has eager_input_streaming: true.
+ {
+ name: "Anthropic/eager_input_streaming_header_added",
+ provider: schemas.Anthropic,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ eager := true
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1", EagerInputStreaming: &eager}},
+ }
+ },
+ expectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
+ {
+ name: "Bedrock/eager_input_streaming_header_added",
+ provider: schemas.Bedrock,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ eager := true
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1", EagerInputStreaming: &eager}},
+ }
+ },
+ expectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
+ {
+ name: "Vertex/eager_input_streaming_header_added",
+ provider: schemas.Vertex,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ eager := true
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1", EagerInputStreaming: &eager}},
+ }
+ },
+ expectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
+ {
+ name: "Azure/eager_input_streaming_header_added",
+ provider: schemas.Azure,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ eager := true
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1", EagerInputStreaming: &eager}},
+ }
+ },
+ expectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
+ {
+ name: "eager_input_streaming_header_skipped_when_flag_false",
+ provider: schemas.Anthropic,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ eager := false
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1", EagerInputStreaming: &eager}},
+ }
+ },
+ unexpectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
+ {
+ name: "eager_input_streaming_header_skipped_when_unset",
+ provider: schemas.Anthropic,
+ setupReq: func() *anthropic.AnthropicMessageRequest {
+ return &anthropic.AnthropicMessageRequest{
+ Tools: []anthropic.AnthropicTool{{Name: "t1"}},
+ }
+ },
+ unexpectHeaders: []string{"fine-grained-tool-streaming-2025-05-14"},
+ },
}
for _, tt := range tests {
diff --git a/core/internal/llmtests/server_tools_via_openai.go b/core/internal/llmtests/server_tools_via_openai.go
new file mode 100644
index 0000000000..c5ee1d2000
--- /dev/null
+++ b/core/internal/llmtests/server_tools_via_openai.go
@@ -0,0 +1,152 @@
+package llmtests
+
+import (
+ "context"
+ "os"
+ "strings"
+ "testing"
+
+ bifrost "github.com/maximhq/bifrost/core"
+ "github.com/maximhq/bifrost/core/schemas"
+)
+
+// RunServerToolsViaOpenAIEndpointTest reproduces the user-reported bug where
+// sending an Anthropic-server-tool-shaped entry in tools[] via the OpenAI-
+// compatible chat-completions endpoint was silently dropped (Claude responded
+// with a prose "I can't check real-time data" fallback). The fix was a
+// combination of:
+// - ChatTool schema gaining Name + all server-tool variant fields.
+// - ToAnthropicChatRequest learning to convert non-function tools (server
+// tools) into AnthropicTool with the correct variant embed.
+//
+// This test sends the exact curl-reported shape via BifrostChatRequest +
+// ChatCompletionRequest and asserts the request succeeds end-to-end against
+// the provider. It covers three server tools that have single-turn triggers
+// (web_search, web_fetch, code_execution) across all supporting providers per
+// Table 20. Other variants (bash, memory, text_editor, tool_search,
+// mcp_toolset, computer_use) require multi-turn tool loops or infra setup
+// and are covered by the schema / unit-level round-trip tests instead.
+func RunServerToolsViaOpenAIEndpointTest(t *testing.T, client *bifrost.Bifrost, ctx context.Context, testConfig ComprehensiveTestConfig) {
+ if !testConfig.Scenarios.ServerToolsViaOpenAIEndpoint {
+ t.Logf("ServerToolsViaOpenAIEndpoint not supported for provider %s", testConfig.Provider)
+ return
+ }
+
+ cases := []struct {
+ name string
+ toolType schemas.ChatToolType
+ toolName string
+ prompt string
+ // extra lets the case set server-tool metadata (max_uses etc.).
+ extra func(*schemas.ChatTool)
+ // supported reports whether this tool is supported on the given
+ // provider per Table 20 (cited provider feature matrix).
+ supported func(schemas.ModelProvider) bool
+ }{
+ {
+ name: "web_search",
+ toolType: "web_search_20260209",
+ toolName: "web_search",
+ prompt: "What is the weather in San Francisco today? Use the web_search tool.",
+ extra: func(t *schemas.ChatTool) {
+ five := 5
+ t.MaxUses = &five
+ t.AllowedCallers = []string{"direct"}
+ },
+ // web_search: Anthropic + Vertex + Azure per Table 20 (not Bedrock).
+ supported: func(p schemas.ModelProvider) bool {
+ return p == schemas.Anthropic || p == schemas.Vertex || p == schemas.Azure
+ },
+ },
+ {
+ name: "web_fetch",
+ toolType: "web_fetch_20260309",
+ toolName: "web_fetch",
+ prompt: "Fetch https://example.com and summarise the title.",
+ extra: func(t *schemas.ChatTool) {
+ three := 3
+ t.MaxUses = &three
+ },
+ // web_fetch: Anthropic + Azure only per Table 20.
+ supported: func(p schemas.ModelProvider) bool {
+ return p == schemas.Anthropic || p == schemas.Azure
+ },
+ },
+ {
+ name: "code_execution",
+ toolType: "code_execution_20250825",
+ toolName: "code_execution",
+ prompt: "Compute 2^64 minus 1 using the code_execution tool and return the result.",
+ // code_execution: Anthropic + Azure only per Table 20.
+ supported: func(p schemas.ModelProvider) bool {
+ return p == schemas.Anthropic || p == schemas.Azure
+ },
+ },
+ }
+
+ t.Run("ServerToolsViaOpenAIEndpoint", func(t *testing.T) {
+ if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
+ t.Parallel()
+ }
+
+ for _, tc := range cases {
+ tc := tc
+ t.Run(tc.name, func(t *testing.T) {
+ if !tc.supported(testConfig.Provider) {
+ t.Skipf("%s not supported on %s per Table 20", tc.name, testConfig.Provider)
+ }
+ if os.Getenv("SKIP_PARALLEL_TESTS") != "true" {
+ t.Parallel()
+ }
+
+ tool := schemas.ChatTool{
+ Type: tc.toolType,
+ Name: tc.toolName,
+ }
+ if tc.extra != nil {
+ tc.extra(&tool)
+ }
+
+ req := &schemas.BifrostChatRequest{
+ Provider: testConfig.Provider,
+ Model: testConfig.ChatModel,
+ Input: []schemas.ChatMessage{
+ CreateBasicChatMessage(tc.prompt),
+ },
+ Params: &schemas.ChatParameters{
+ MaxCompletionTokens: bifrost.Ptr(500),
+ Tools: []schemas.ChatTool{tool},
+ },
+ Fallbacks: testConfig.Fallbacks,
+ }
+
+ bfCtx := schemas.NewBifrostContext(ctx, schemas.NoDeadline)
+ resp, err := client.ChatCompletionRequest(bfCtx, req)
+ if err != nil {
+ t.Fatalf("%s tool request failed: %s", tc.name, GetErrorMessage(err))
+ }
+ if resp == nil {
+ t.Fatal("expected non-nil response")
+ }
+
+ // Regression signals:
+ // 1. Upstream accepted the request (no error).
+ // 2. Response is not the prose fallback Claude emits when
+ // the server-tool was silently stripped pre-fix
+ // ("I can't/cannot/don't have access to real-time ...").
+ // The schema + conversion unit tests prove the outbound
+ // request carries the tool; this live test proves the
+ // provider accepts the shape AND actually uses the tool
+ // rather than answering from parametric memory.
+ content := GetChatContent(resp)
+ lc := strings.ToLower(content)
+ if strings.Contains(lc, "can't access real-time") ||
+ strings.Contains(lc, "cannot access real-time") ||
+ strings.Contains(lc, "don't have access to real-time") {
+ t.Fatalf("%s regression: tool appears to be ignored, content=%q", tc.name, content)
+ }
+ t.Logf("%s tool live call succeeded: chars=%d", tc.name, len(content))
+ })
+ }
+ })
+}
diff --git a/core/internal/llmtests/tests.go b/core/internal/llmtests/tests.go
index af3006b9a1..108894feb4 100644
--- a/core/internal/llmtests/tests.go
+++ b/core/internal/llmtests/tests.go
@@ -120,6 +120,8 @@ func RunAllComprehensiveTests(t *testing.T, client *bifrost.Bifrost, ctx context
RunCompactionTest,
RunInterleavedThinkingTest,
RunFastModeTest,
+ RunEagerInputStreamingTest,
+ RunServerToolsViaOpenAIEndpointTest,
}
// Execute all test scenarios without raw request/response (default behavior)
@@ -239,6 +241,8 @@ func printTestSummary(t *testing.T, testConfig ComprehensiveTestConfig) {
{"Compaction", testConfig.Scenarios.Compaction},
{"InterleavedThinking", testConfig.Scenarios.InterleavedThinking},
{"FastMode", testConfig.Scenarios.FastMode},
+ {"EagerInputStreaming", testConfig.Scenarios.EagerInputStreaming},
+ {"ServerToolsViaOpenAIEndpoint", testConfig.Scenarios.ServerToolsViaOpenAIEndpoint},
}
supported := 0
diff --git a/core/internal/mcptests/annotations_test.go b/core/internal/mcptests/annotations_test.go
new file mode 100644
index 0000000000..e85b54a79f
--- /dev/null
+++ b/core/internal/mcptests/annotations_test.go
@@ -0,0 +1,220 @@
+package mcptests
+
+import (
+ "encoding/json"
+ "strings"
+ "testing"
+
+ "github.com/maximhq/bifrost/core/schemas"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+// =============================================================================
+// MCP ANNOTATION TESTS
+//
+// These tests verify two invariants of the MCP annotations feature:
+//
+// 1. PRESERVATION: annotations attached to a registered tool survive the full
+// MCP→Bifrost conversion and remain accessible on ChatTool.Annotations
+// after retrieval from the manager.
+//
+// 2. ISOLATION: annotations are tagged json:"-" on ChatTool, so they are never
+// included in the JSON body forwarded to LLM providers.
+// =============================================================================
+
+// TestAnnotations_PreservedAfterToolRegistration verifies that annotations set
+// on an InProcess ChatTool schema are stored in the tool map without modification.
+func TestAnnotations_PreservedAfterToolRegistration(t *testing.T) {
+ t.Parallel()
+
+ readOnly := true
+ idempotent := true
+
+ manager := setupMCPManager(t)
+
+ toolSchema := schemas.ChatTool{
+ Type: schemas.ChatToolTypeFunction,
+ Function: &schemas.ChatToolFunction{
+ Name: "read_resource",
+ Description: schemas.Ptr("Reads a resource"),
+ Parameters: &schemas.ToolFunctionParameters{
+ Type: "object",
+ Properties: schemas.NewOrderedMapFromPairs(
+ schemas.KV("uri", map[string]interface{}{
+ "type": "string",
+ "description": "URI of the resource to read",
+ }),
+ ),
+ Required: []string{"uri"},
+ },
+ },
+ Annotations: &schemas.MCPToolAnnotations{
+ Title: "Resource Reader",
+ ReadOnlyHint: &readOnly,
+ IdempotentHint: &idempotent,
+ },
+ }
+
+ err := manager.RegisterTool(
+ "read_resource",
+ "Reads a resource",
+ func(args any) (string, error) { return `{"ok":true}`, nil },
+ toolSchema,
+ )
+ require.NoError(t, err)
+
+ ctx := createTestContext()
+ toolPerClient := manager.GetToolPerClient(ctx)
+
+ var found *schemas.ChatTool
+outer1:
+ for _, tools := range toolPerClient {
+ for i := range tools {
+ if tools[i].Function != nil && strings.HasSuffix(tools[i].Function.Name, "-read_resource") {
+ cp := tools[i]
+ found = &cp
+ break outer1
+ }
+ }
+ }
+ require.NotNil(t, found, "read_resource tool should be present in the tool map")
+
+ // Annotations must be preserved on ChatTool (not lost after registration)
+ require.NotNil(t, found.Annotations, "Annotations should be preserved on ChatTool")
+ assert.Equal(t, "Resource Reader", found.Annotations.Title)
+ require.NotNil(t, found.Annotations.ReadOnlyHint)
+ assert.True(t, *found.Annotations.ReadOnlyHint)
+ require.NotNil(t, found.Annotations.IdempotentHint)
+ assert.True(t, *found.Annotations.IdempotentHint)
+ assert.Nil(t, found.Annotations.DestructiveHint)
+ assert.Nil(t, found.Annotations.OpenWorldHint)
+}
+
+// TestAnnotations_AbsentFromProviderJSON verifies that annotations do NOT appear
+// in the JSON representation of a tool — i.e. the payload that would be forwarded
+// to an LLM provider.
+func TestAnnotations_AbsentFromProviderJSON(t *testing.T) {
+ t.Parallel()
+
+ readOnly := true
+ destructive := false
+
+ manager := setupMCPManager(t)
+
+ toolSchema := schemas.ChatTool{
+ Type: schemas.ChatToolTypeFunction,
+ Function: &schemas.ChatToolFunction{
+ Name: "write_file",
+ Description: schemas.Ptr("Writes content to a file"),
+ Parameters: &schemas.ToolFunctionParameters{
+ Type: "object",
+ Properties: schemas.NewOrderedMapFromPairs(
+ schemas.KV("path", map[string]interface{}{
+ "type": "string",
+ "description": "Destination file path",
+ }),
+ schemas.KV("content", map[string]interface{}{
+ "type": "string",
+ "description": "Content to write",
+ }),
+ ),
+ Required: []string{"path", "content"},
+ },
+ },
+ Annotations: &schemas.MCPToolAnnotations{
+ Title: "File Writer",
+ ReadOnlyHint: &readOnly,
+ DestructiveHint: &destructive,
+ },
+ }
+
+ err := manager.RegisterTool(
+ "write_file",
+ "Writes content to a file",
+ func(args any) (string, error) { return `{"ok":true}`, nil },
+ toolSchema,
+ )
+ require.NoError(t, err)
+
+ ctx := createTestContext()
+ toolPerClient := manager.GetToolPerClient(ctx)
+
+ var found *schemas.ChatTool
+outer2:
+ for _, tools := range toolPerClient {
+ for i := range tools {
+ if tools[i].Function != nil && strings.HasSuffix(tools[i].Function.Name, "-write_file") {
+ cp := tools[i]
+ found = &cp
+ break outer2
+ }
+ }
+ }
+ require.NotNil(t, found, "write_file tool should be present in the tool map")
+
+ // The tool must have annotations in memory
+ require.NotNil(t, found.Annotations, "Annotations must be in memory for downstream use")
+
+ // Serialize the tool as a provider would receive it
+ toolJSON, err := json.Marshal(found)
+ require.NoError(t, err)
+ s := string(toolJSON)
+
+ // None of the annotation data must leak into the JSON.
+ // Use the key token `"annotations":` to avoid false positives from description text.
+ assert.NotContains(t, s, `"annotations":`, "annotations key must be absent from provider JSON")
+ assert.NotContains(t, s, "readOnlyHint", "readOnlyHint must be absent from provider JSON")
+ assert.NotContains(t, s, "destructiveHint", "destructiveHint must be absent from provider JSON")
+ assert.NotContains(t, s, "File Writer", "annotation title must be absent from provider JSON")
+
+ // The function definition itself must still be present
+ assert.Contains(t, s, "write_file", "function name must be present in provider JSON")
+ assert.Contains(t, s, "path", "parameter must be present in provider JSON")
+}
+
+// TestAnnotations_DeepCopyPreservesAnnotations verifies that the deep-copy path
+// (used during plugin accumulation and streaming) correctly copies annotations.
+func TestAnnotations_DeepCopyPreservesAnnotations(t *testing.T) {
+ t.Parallel()
+
+ readOnly := true
+
+ original := schemas.ChatTool{
+ Type: schemas.ChatToolTypeFunction,
+ Function: &schemas.ChatToolFunction{
+ Name: "read_config",
+ Description: schemas.Ptr("Reads configuration from disk"),
+ },
+ Annotations: &schemas.MCPToolAnnotations{
+ Title: "Config Reader",
+ ReadOnlyHint: &readOnly,
+ },
+ }
+
+ copied := schemas.DeepCopyChatTool(original)
+
+ // Annotations must survive the deep copy
+ require.NotNil(t, copied.Annotations, "Annotations must be preserved after deep copy")
+ assert.Equal(t, "Config Reader", copied.Annotations.Title)
+ require.NotNil(t, copied.Annotations.ReadOnlyHint)
+ assert.True(t, *copied.Annotations.ReadOnlyHint)
+
+ // Mutate via the pointed-to value to detect pointer aliasing
+ *original.Annotations.ReadOnlyHint = false
+ assert.NotSame(t, original.Annotations.ReadOnlyHint, copied.Annotations.ReadOnlyHint,
+ "deep copy must not share the ReadOnlyHint pointer with the original")
+ assert.True(t, *copied.Annotations.ReadOnlyHint,
+ "mutating original's ReadOnlyHint must not affect the deep copy")
+
+ // JSON of the copy must also be annotation-free (same guarantee as the original)
+ toolJSON, err := json.Marshal(copied)
+ require.NoError(t, err)
+ s := string(toolJSON)
+ // Check for the JSON key pattern, not just the substring, to avoid false positives
+ // from description text. The key would appear as `"annotations":` in JSON.
+ assert.NotContains(t, s, `"annotations":`,
+ "annotations key must be absent from provider JSON even after deep copy")
+ assert.NotContains(t, s, "readOnlyHint",
+ "readOnlyHint must be absent from provider JSON even after deep copy")
+}
diff --git a/core/mcp/utils.go b/core/mcp/utils.go
index d80ec17acc..1356bb38bb 100644
--- a/core/mcp/utils.go
+++ b/core/mcp/utils.go
@@ -487,6 +487,28 @@ func convertMCPToolToBifrostSchema(mcpTool *mcp.Tool, logger schemas.Logger) sch
// object schemas to always have a properties field, even if empty
properties = schemas.NewOrderedMap()
}
+
+ // Preserve MCP tool annotations if any are set.
+ // Clone bool pointers so Bifrost's copy is independent of the upstream mcp.Tool lifetime.
+ var annotations *schemas.MCPToolAnnotations
+ a := mcpTool.Annotations
+ if a.Title != "" || a.ReadOnlyHint != nil || a.DestructiveHint != nil || a.IdempotentHint != nil || a.OpenWorldHint != nil {
+ cloneBool := func(b *bool) *bool {
+ if b == nil {
+ return nil
+ }
+ v := *b
+ return &v
+ }
+ annotations = &schemas.MCPToolAnnotations{
+ Title: a.Title,
+ ReadOnlyHint: cloneBool(a.ReadOnlyHint),
+ DestructiveHint: cloneBool(a.DestructiveHint),
+ IdempotentHint: cloneBool(a.IdempotentHint),
+ OpenWorldHint: cloneBool(a.OpenWorldHint),
+ }
+ }
+
return schemas.ChatTool{
Type: schemas.ChatToolTypeFunction,
Function: &schemas.ChatToolFunction{
@@ -498,6 +520,7 @@ func convertMCPToolToBifrostSchema(mcpTool *mcp.Tool, logger schemas.Logger) sch
Required: mcpTool.InputSchema.Required,
},
},
+ Annotations: annotations,
}
}
diff --git a/core/mcp/utils_test.go b/core/mcp/utils_test.go
index e74d9d7da1..1fba67db38 100644
--- a/core/mcp/utils_test.go
+++ b/core/mcp/utils_test.go
@@ -1,9 +1,13 @@
package mcp
import (
+ "encoding/json"
"testing"
"github.com/mark3labs/mcp-go/mcp"
+ "github.com/maximhq/bifrost/core/schemas"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
)
// TestConvertMCPToolToBifrostSchema_EmptyParameters tests that tools with no parameters
@@ -49,6 +53,68 @@ func TestConvertMCPToolToBifrostSchema_EmptyParameters(t *testing.T) {
}
}
+// TestConvertMCPToolToBifrostSchema_WithAnnotations tests that MCP tool annotations
+// are preserved on ChatTool.Annotations (not ChatToolFunction) and are absent from JSON.
+func TestConvertMCPToolToBifrostSchema_WithAnnotations(t *testing.T) {
+ readOnly := true
+ destructive := false
+
+ mcpTool := &mcp.Tool{
+ Name: "read_resource",
+ Description: "Reads a resource",
+ InputSchema: mcp.ToolInputSchema{
+ Type: "object",
+ Properties: map[string]interface{}{},
+ },
+ Annotations: mcp.ToolAnnotation{
+ Title: "Resource Reader",
+ ReadOnlyHint: &readOnly,
+ DestructiveHint: &destructive,
+ IdempotentHint: schemas.Ptr(true),
+ },
+ }
+
+ bifrostTool := convertMCPToolToBifrostSchema(mcpTool, defaultLogger)
+
+ // Annotations must be on ChatTool, not buried in Function
+ require.NotNil(t, bifrostTool.Annotations, "Annotations should be set on ChatTool")
+ assert.Equal(t, "Resource Reader", bifrostTool.Annotations.Title)
+ require.NotNil(t, bifrostTool.Annotations.ReadOnlyHint)
+ assert.True(t, *bifrostTool.Annotations.ReadOnlyHint)
+ require.NotNil(t, bifrostTool.Annotations.DestructiveHint)
+ assert.False(t, *bifrostTool.Annotations.DestructiveHint)
+ require.NotNil(t, bifrostTool.Annotations.IdempotentHint)
+ assert.True(t, *bifrostTool.Annotations.IdempotentHint)
+ assert.Nil(t, bifrostTool.Annotations.OpenWorldHint)
+
+ // The JSON sent to providers must not contain annotations
+ toolJSON, err := json.Marshal(bifrostTool)
+ require.NoError(t, err)
+ s := string(toolJSON)
+ assert.NotContains(t, s, "annotations", "annotations must be absent from provider JSON")
+ assert.NotContains(t, s, "readOnlyHint", "readOnlyHint must be absent from provider JSON")
+ assert.NotContains(t, s, "Resource Reader", "annotation title must be absent from provider JSON")
+}
+
+// TestConvertMCPToolToBifrostSchema_NilAnnotationsWhenAllZero verifies the nil guard:
+// when all annotation fields are zero-valued, ChatTool.Annotations must remain nil.
+func TestConvertMCPToolToBifrostSchema_NilAnnotationsWhenAllZero(t *testing.T) {
+ mcpTool := &mcp.Tool{
+ Name: "no_hints_tool",
+ Description: "A tool with no annotation hints",
+ InputSchema: mcp.ToolInputSchema{
+ Type: "object",
+ Properties: map[string]interface{}{},
+ },
+ Annotations: mcp.ToolAnnotation{}, // All zero values — Title empty, all hints nil
+ }
+
+ bifrostTool := convertMCPToolToBifrostSchema(mcpTool, defaultLogger)
+
+ assert.Nil(t, bifrostTool.Annotations,
+ "Annotations should be nil when all MCP annotation fields are zero")
+}
+
// TestConvertMCPToolToBifrostSchema_WithParameters tests the normal case with parameters
func TestConvertMCPToolToBifrostSchema_WithParameters(t *testing.T) {
// Create a tool with parameters
diff --git a/core/providers/anthropic/anthropic.go b/core/providers/anthropic/anthropic.go
index e5d62fe87c..e012f3a13f 100644
--- a/core/providers/anthropic/anthropic.go
+++ b/core/providers/anthropic/anthropic.go
@@ -450,6 +450,28 @@ func (provider *AnthropicProvider) ChatCompletion(ctx *schemas.BifrostContext, k
return nil, bifrostErr
}
+ // On the raw-body passthrough path, the typed-struct StripUnsupportedAnthropicFields
+ // was not invoked. Apply the JSON-level sanitizer for behavioural parity so
+ // unsupported request-level and tool-level fields don't leak to providers that
+ // would reject them.
+ if useRawBody, ok := ctx.Value(schemas.BifrostContextKeyUseRawRequestBody).(bool); ok && useRawBody {
+ // Feature gating keyed to schemas.Anthropic (not provider.GetProviderKey())
+ // so custom Anthropic aliases get the same feature lookup as the typed
+ // path above (line 445), keeping raw and typed behavior in lockstep.
+ sanitized, rawErr := stripUnsupportedFieldsFromRawBody(jsonData, schemas.Anthropic, request.Model)
+ if rawErr != nil {
+ return nil, providerUtils.NewBifrostOperationError(schemas.ErrProviderRequestMarshal, rawErr, provider.GetProviderKey())
+ }
+ jsonData = sanitized
+ // Auto-inject matching anthropic-beta headers for fields the sanitizer
+ // preserved. Probe-unmarshal reuses the typed path's header walker so
+ // the two paths stay in lockstep.
+ var probe AnthropicMessageRequest
+ if err := schemas.Unmarshal(jsonData, &probe); err == nil {
+ AddMissingBetaHeadersToContext(ctx, &probe, schemas.Anthropic)
+ }
+ }
+
// Use struct directly for JSON marshaling
responseBody, latency, providerResponseHeaders, err := provider.completeRequest(ctx, jsonData, provider.buildRequestURL(ctx, "/v1/messages", schemas.ChatCompletionRequest), key.Value.GetValue(), &providerUtils.RequestMetadata{
Provider: provider.GetProviderKey(),
@@ -534,6 +556,25 @@ func (provider *AnthropicProvider) ChatCompletionStream(ctx *schemas.BifrostCont
return nil, bifrostErr
}
+ // On the raw-body passthrough path, the typed-struct StripUnsupportedAnthropicFields
+ // was not invoked. Apply the JSON-level sanitizer for behavioural parity.
+ if useRawBody, ok := ctx.Value(schemas.BifrostContextKeyUseRawRequestBody).(bool); ok && useRawBody {
+ // Feature gating keyed to schemas.Anthropic (not provider.GetProviderKey())
+ // to keep raw and typed paths in lockstep on custom aliases — mirrors
+ // the typed path's hardcoded schemas.Anthropic at line 548.
+ sanitized, rawErr := stripUnsupportedFieldsFromRawBody(jsonData, schemas.Anthropic, request.Model)
+ if rawErr != nil {
+ return nil, providerUtils.NewBifrostOperationError(schemas.ErrProviderRequestMarshal, rawErr, provider.GetProviderKey())
+ }
+ jsonData = sanitized
+ // Auto-inject matching anthropic-beta headers for fields the sanitizer
+ // preserved. Probe-unmarshal reuses the typed path's header walker.
+ var probe AnthropicMessageRequest
+ if err := schemas.Unmarshal(jsonData, &probe); err == nil {
+ AddMissingBetaHeadersToContext(ctx, &probe, schemas.Anthropic)
+ }
+ }
+
// Prepare Anthropic headers
headers := map[string]string{
"Content-Type": "application/json",
@@ -660,6 +701,7 @@ func HandleAnthropicChatCompletionStreaming(
// Start streaming in a goroutine
go func() {
+ defer providerUtils.EnsureStreamFinalizerCalled(ctx)
defer func() {
model := "unknown"
if meta != nil {
@@ -1146,6 +1188,7 @@ func HandleAnthropicResponsesStream(
// Start streaming in a goroutine
go func() {
+ defer providerUtils.EnsureStreamFinalizerCalled(ctx)
defer func() {
model := "
+
+
+You will need to make 2 changes in the app manifest
+
+```json
+"requestedAccessTokenVersion": 2
+```
+
+and
+
+```json
+"optionalClaims": {
+ "idToken": [
+ {
+ "name": "roles",
+ "source": null,
+ "essential": false,
+ "additionalProperties": []
+ }
+ ],
+ "accessToken": [],
+ "saml2Token": []
+}
+```
+
+## Step 9: Configure Bifrost
Now configure Bifrost to use Microsoft Entra as the identity provider.
### Using the Bifrost UI
+
+ 
+
+
1. Navigate to **Governance** → **User Provisioning** in your Bifrost dashboard
2. Select **Microsoft Entra** as the SCIM Provider
3. Enter the following configuration:
@@ -221,8 +254,9 @@ Now configure Bifrost to use Microsoft Entra as the identity provider.
| **Audience** | Your Client ID (optional, defaults to Client ID) |
| **App ID URI** | `api://{client-id}` (optional, for v1.0 tokens) |
-4. Toggle **Enabled** to activate the provider
-5. Click **Save Configuration**
+5. **Verify** configuration and see if you get any errors. Make sure you get no errors/warnings.
+6. Toggle **Enabled** to activate the provider
+7. Click **Save Configuration**
+
+
+
+
+
+
+#### Evaluation rules
+
+- **Role mappings**: Ordered, first match wins. If no rule matches, users are not allowed to login into the system.
+- **Team and business unit mappings**: All matching rules apply — users can be placed on multiple teams and business units simultaneously.
+- **Claim values**: Can be strings, arrays, or nested objects. Bifrost resolves dotted paths (e.g., `realm_access.roles`).
+
+
+
+
+2. Choose **Internal** if you only want Workspace users, or **External** otherwise.
+3. Fill in App name, support email, and developer contact.
+4. Add the scopes: `openid`, `profile`, `email`.
+5. Save.
+
+---
+
+## Step 2: Create an OAuth Client ID
+
+1. Open **APIs & Services → Credentials → Create credentials → OAuth client ID**.
+
+
+ 
+
+
+2. Configure:
+
+| Field | Value |
+| --- | --- |
+| **Application type** | Web application |
+| **Name** | Bifrost Enterprise |
+| **Authorized JavaScript origins** | `https://your-bifrost-domain.com` |
+| **Authorized redirect URIs** | `https://your-bifrost-domain.com/login` |
+
+3. Save and copy the **Client ID** and **Client Secret**.
+
+---
+
+## Step 3: (Optional) Create a service account for Directory API access
+
+Skip this section if you only want SSO login without directory-based user import.
+
+1. Go to **IAM & Admin → Service Accounts → Create service account**.
+
+
+ 
+
+
+2. Give it a name (e.g. `bifrost-provisioning`). You can skip the "Grant this service account access to project" step — no GCP IAM roles are required; access is granted via domain-wide delegation in Step 5.
+3. Open the service account → **Keys → Add Key → Create new key → JSON**. Download and store the JSON file securely.
+4. From the service account **Details** tab, copy the **Unique ID** (a numeric value, **not** the email or OAuth Client ID).
+
+---
+
+## Step 4: Enable the Admin SDK API
+
+If you're using the service account path:
+
+1. Open **APIs & Services → Library**.
+2. Search for **Admin SDK API** and click **Enable**.
+
+---
+
+## Step 5: Set up domain-wide delegation
+
+1. In the [Google Admin Console](https://admin.google.com), go to **Security → Access and data control → API controls → Manage Domain Wide Delegation**.
+
+
+ 
+
+
+2. Click **Add new**.
+3. Enter the service account's **Unique ID** (from Step 3).
+4. Add these OAuth scopes (copy the full URLs, comma-separated):
+
+```
+https://www.googleapis.com/auth/admin.directory.user.readonly,
+https://www.googleapis.com/auth/admin.directory.group.readonly,
+https://www.googleapis.com/auth/admin.directory.group.member.readonly
+```
+
+5. **Authorize**.
+
+
+
+
+4. Click **Verify** — Bifrost validates the OAuth client and, if a service account is provided, attempts a Directory API impersonation to confirm delegation is working.
+5. Configure **Attribute → Role / Team / Business Unit** mappings to map groups or organizational units to Bifrost roles and teams.
+6. Toggle **Enabled** and click **Save Configuration**.
+
+### Using `config.json`
+
+```json
+{
+ "scim_config": {
+ "enabled": true,
+ "provider": "google",
+ "config": {
+ "domain": "company.com",
+ "clientId": "123-abc.apps.googleusercontent.com",
+ "clientSecret": "${GOOGLE_WORKSPACE_CLIENT_SECRET}",
+ "adminEmail": "sso-admin@company.com",
+ "serviceAccountEnvVar": "GOOGLE_SA_JSON",
+ "teamIdsField": "groups"
+ }
+ }
+}
+```
+
+Pick one of the three service-account sources: `serviceAccountJson` (raw JSON string), `serviceAccountEnvVar` (env var name holding the JSON), or `serviceAccountFile` (absolute path to the key file).
+
+
+### Custom attribute mapping
+
+You can also map any custom attributes to any entity (role, team or business unit). Make sure these are configured to send back to Bifrost in token configuration.
+
+
+
+
+
+### Configuration reference
+
+| Field | Required | Description |
+| --- | --- | --- |
+| `domain` | Yes | Google Workspace primary domain (e.g. `company.com`). |
+| `clientId` | Yes | OAuth 2.0 Web Client ID from Step 2. |
+| `clientSecret` | Yes | Client Secret — required for token revocation and for confidential server-side flows. |
+| `audience` | No | Expected JWT audience. Defaults to `clientId`. |
+| `adminEmail` | Yes | Workspace admin to impersonate via domain-wide delegation. Required when any service-account field is set. |
+| `serviceAccountJson` | One of 3 | Raw JSON string of the service account key. |
+| `serviceAccountEnvVar` | One of 3 | Name of the environment variable containing the JSON. |
+| `serviceAccountFile` | One of 3 | Absolute path to the JSON key file on the Bifrost host. |
+| `attributeRoleMappings` | Yes | Ordered list of attribute→role mappings. |
+| `attributeTeamMappings` | No | Attribute→team mappings (all matches apply). |
+| `attributeBusinessUnitMappings` | No | Attribute→business-unit mappings (all matches apply). |
+
+
+
+
+2. Configure the client:
+
+| Field | Value |
+| --- | --- |
+| **Client type** | OpenID Connect |
+| **Client ID** | `bifrost` (or your preferred identifier) |
+| **Name** | `Bifrost Enterprise` |
+
+3. On the **Capability config** step enable:
+ - **Client authentication** (makes it a confidential client)
+ - **Standard flow** (Authorization Code)
+ - **Service accounts roles** (required for Admin REST API access)
+
+
+ 
+
+
+4. On the **Login settings** step set:
+
+| Field | Value |
+| --- | --- |
+| **Valid redirect URIs** | `https://your-bifrost-domain.com/login` |
+| **Valid post logout redirect URIs** | `https://your-bifrost-domain.com` |
+| **Web origins** | `https://your-bifrost-domain.com` |
+
+5. **Save** the client.
+
+---
+
+## Step 2: Copy the client credentials
+
+1. Open the client and go to the **Credentials** tab.
+2. Copy the **Client Secret**.
+
+
+ 
+
+
+---
+
+## Step 3: Configure role and group mappers
+
+Keycloak does not include realm roles or full group paths in tokens by default. Add two mappers on the client's dedicated scope.
+
+1. Open the client → **Client Scopes** tab → click the client's `-dedicated` scope.
+2. Click **Add mapper → By configuration**.
+
+### Group Membership mapper
+
+
+ 
+
+
+| Field | Value |
+| --- | --- |
+| **Mapper Type** | Group Membership |
+| **Name** | `groups` |
+| **Token Claim Name** | `groups` |
+| **Full group path** | **On** |
+| **Add to ID token** | **On** |
+| **Add to access token** | **On** |
+| **Add to userinfo** | **On** |
+
+
+
+
+---
+
+## Step 5: Create realm roles and groups
+
+Create the roles and groups you plan to map into Bifrost.
+
+1. **Realm → Realm roles → Create role** for each role (e.g. `bifrost-admin`, `bifrost-developer`, `bifrost-viewer`).
+2. **Realm → Groups → Create group** for each team you want to sync (e.g. `/platform`, `/data-science`).
+3. Assign users to the appropriate roles and groups under **Users → your user → Role mapping** / **Groups**.
+
+---
+
+## Step 6: Configure Bifrost
+
+### Using the Bifrost dashboard
+
+1. In Bifrost, go to **Governance → User Provisioning**.
+2. Select **Keycloak** as the SCIM Provider.
+3. Fill in the fields:
+
+| Field | Value |
+| --- | --- |
+| **Server URL** | `https://keycloak.company.com` (no `/realms/...` suffix) |
+| **Realm** | Your realm name (e.g. `master`, `bifrost-prod`) |
+| **Client ID** | Client ID from Step 1 |
+| **Client Secret** | Client Secret from Step 2 |
+| **Audience** | Optional — defaults to Client ID |
+| **Team IDs Field** | Leave as `groups` (default) or change if you used a different mapper name |
+
+4. Click **Verify** — Bifrost connects to Keycloak's JWKS and Admin REST API to confirm the client and service-account roles.
+5. Configure **Attribute → Role / Team / Business Unit** mappings if needed.
+6. Toggle **Enabled** and click **Save Configuration**.
+
+
+ 
+
+
+### Using `config.json`
+
+```json
+{
+ "scim_config": {
+ "enabled": true,
+ "provider": "keycloak",
+ "config": {
+ "serverUrl": "https://keycloak.company.com",
+ "realm": "bifrost-prod",
+ "clientId": "bifrost",
+ "clientSecret": "${KEYCLOAK_CLIENT_SECRET}",
+ "teamIdsField": "groups"
+ }
+ }
+}
+```
+
+### Configuration reference
+
+| Field | Required | Description |
+| --- | --- | --- |
+| `serverUrl` | Yes | Base URL of the Keycloak server. Must be a valid URL (e.g. `https://keycloak.company.com`) and must **not** include `/realms/...`. |
+| `realm` | Yes | Realm name. |
+| `clientId` | Yes | Client ID created in Step 1. |
+| `clientSecret` | Yes | Client secret — required because the client is confidential. |
+| `audience` | No | Expected JWT audience. Defaults to `clientId`. |
+| `teamIdsField` | No | JWT claim for group IDs. Defaults to `groups`. |
+| `attributeRoleMappings` | No | Ordered list of attribute→role mappings. |
+| `attributeTeamMappings` | No | Attribute→team mappings (all matches apply). |
+| `attributeBusinessUnitMappings` | No | Attribute→business-unit mappings (all matches apply). |
+
+---
+
+## Testing the Integration
+
+1. Open the Bifrost dashboard in an incognito window.
+2. You'll be redirected to Keycloak's login page.
+3. Sign in with a Keycloak user that has one of the roles you mapped.
+4. Verify the user appears under **Governance → Users** with the expected role and teams.
+5. From **Governance → User Provisioning → Import Users**, verify the service account can list users.
+
+---
+
+## Troubleshooting
+
+### `serverUrl must not include /realms/{realm}`
+
+The `serverUrl` field is the base Keycloak URL. Set the realm in the separate **Realm** field. Example: `https://keycloak.company.com` + realm `bifrost-prod` — **not** `https://keycloak.company.com/realms/bifrost-prod`.
+
+### Users redirected back to login
+
+- Confirm the client's **Valid redirect URIs** exactly match your Bifrost login URL (trailing slash matters).
+- Verify the client is **Enabled** in Keycloak.
+
+### Roles not appearing in the token
+
+- Check that the **User Realm Role** mapper adds to both ID and Access tokens.
+- Use `Evaluate` on the client scope to preview the token a user would receive.
+
+### Service account cannot list users
+
+- Confirm `realm-management → view-users` is assigned under **Service accounts roles**.
+- If you enabled **Authorization** on the client, service account tokens may not work — disable Authorization (fine-grained authz) for this client.
+
+### `jwks keys not found`
+
+- Make sure the server URL is reachable from Bifrost. The JWKS endpoint is `{serverUrl}/realms/{realm}/protocol/openid-connect/certs`.
+
+---
+
+## Next Steps
+
+- [User Provisioning overview](./user-provisioning) — capabilities, attribute mappings, bulk import
+- [Role-Based Access Control](./rbac) — configure custom roles before mapping
+- [Audit Logs](./audit-logs) — track authentication events
diff --git a/docs/enterprise/setting-up-okta.mdx b/docs/enterprise/setting-up-okta.mdx
index 7b6f25f0bc..26ac5e8225 100644
--- a/docs/enterprise/setting-up-okta.mdx
+++ b/docs/enterprise/setting-up-okta.mdx
@@ -13,6 +13,7 @@ This guide walks you through configuring Okta as your identity provider for Bifr
- An Okta organization with admin access
- Bifrost Enterprise deployed and accessible
- The redirect URI for your Bifrost instance (e.g., `https://your-bifrost-domain.com/login`)
+- Ensure you have created all the [roles in Bifrost](/enterprise/rbac) that you are aiming to map to with Okta.
---
@@ -47,20 +48,25 @@ Configure the following settings for your application:
**General Settings:**
+
- **App integration name**: `Bifrost Enterprise`
- **Logo** (optional): You can upload the Bifrost logo from [https://www.getmaxim.ai/bifrost/bifrost-logo-only.png](https://www.getmaxim.ai/bifrost/bifrost-logo-only.png)
**Grant type:**
+
- Enable **Authorization Code**
- Enable **Refresh Token**
**Sign-in redirect URIs:**
+
- Add your Bifrost login callback URL: `https://your-bifrost-domain.com/login`
**Sign-out redirect URIs (Optional):**
+
- Add your Bifrost base URL: `https://your-bifrost-domain.com`
**Assignments:**
+
- Choose **Skip group assignment for now** (we'll configure this later)
6. Click **Save** to create the application
@@ -71,39 +77,13 @@ Configure the following settings for your application:
---
-## Step 3: Configure Authorization Server (optional)
+## Step 3: Create Custom Role Attribute (Optional)
-
-
-3. Note the **Issuer URI** for your authorization server (e.g., `https://your-domain.okta.com/oauth2/default`)
-
-
-| Field | Value |
-|-------|-------|
-| **Data type** | string |
-| **Display name** | bifrostRole |
-| **Variable name** | bifrostRole |
-| **Enum** | Check "Define enumerated list of values" |
+| Field | Value |
+| --------------------- | ----------------------------------------------------------- |
+| **Data type** | string |
+| **Display name** | bifrostRole |
+| **Variable name** | bifrostRole |
+| **Enum** | Check "Define enumerated list of values" |
| **Attribute members** | Admin → `admin`, Developer → `developer`, Viewer → `viewer` |
-| **Attribute type** | Personal |
+| **Attribute type** | Personal |
5. Click **Save**
---
-## Step 5: Add Role Claim to Tokens
+## Step 4: Add Role Claim to Tokens (If you have added custom role attribute)
Configure the authorization server to include the role in the access token.
@@ -148,27 +128,29 @@ Configure the authorization server to include the role in the access token.
Configure the claim:
-| Field | Value |
-|-------|-------|
-| **Name** | `role` |
+| Field | Value |
+| ------------------------- | -------------------- |
+| **Name** | `role` |
| **Include in token type** | Access Token, Always |
-| **Value type** | Expression |
-| **Value** | `user.bifrostRole` |
-| **Include in** | Any scope |
+| **Value type** | Expression |
+| **Value** | `user.bifrostRole` |
+| **Include in** | Any scope |
5. Click **Create**
-
-
-5. Click **Save**
-
-
-
-
-2. Under **OpenID Connect ID Token**, configure:
- - **Groups claim type**: Expression
- - **Groups claim expression**: `Arrays.flatten(Groups.startsWith("OKTA", "bifrost", 100))`
-
-
@@ -271,72 +215,121 @@ Adjust the group filter expression based on your naming convention. The example
4. Click **Save and Go Back**
-
+
+
+1. Click on "Create token"
+
+
+ 
+
+
+2. Copy token to be used in the next step.
+
## Step 8: Configure Bifrost
Now configure Bifrost to use Okta as the identity provider.
### Using the Bifrost UI
+
+ 
+
+
1. Navigate to **Governance** → **User Provisioning** in your Bifrost dashboard
2. Select **Okta** as the SCIM Provider
3. Enter the following configuration:
-| Field | Value |
-|-------|-------|
-| **Client ID** | Your Okta application Client ID |
-| **Issuer URL** | Issuer URL |
-| **Audience** | Your API audience (e.g., `api://default` or custom) |
+| Field | Value |
+| ----------------- | -------------------------------------------------------------------- |
+| **Client ID** | Your Okta application Client ID |
+| **Issuer URL** | Issuer URL |
+| **Audience** | Your API audience (e.g., `api://default` or custom) |
| **Client Secret** | Your Okta application Client Secret (optional, for token revocation) |
-4. Toggle **Enabled** to activate the provider
-5. Click **Save Configuration**
+4. **Verify** configuration and see if you get any errors. Make sure you get no errors/warnings.
+5. Toggle **Enabled** to activate the provider
+6. Click **Save Configuration**
+
+
+
-| Group Claim Name | Role |
-|------------------|------|
-| `bifrost-staging-admins` | Admin |
-| `bifrost-staging-viewers` | Viewer |
+
+
+
+
+#### Evaluation rules
+
+- **Role mappings**: Ordered, first match wins. If no rule matches, users are not allowed to login into the system.
+- **Team and business unit mappings**: All matching rules apply — users can be placed on multiple teams and business units simultaneously.
+- **Claim values**: Can be strings, arrays, or nested objects. Bifrost resolves dotted paths (e.g., `realm_access.roles`).
+
+
+
+2. Configure the application:
+
+| Field | Value |
+| --- | --- |
+| **Type** | Web |
+| **Authentication method** | PKCE (recommended) or Basic |
+| **Redirect URI** | `https://your-bifrost-domain.com/login` |
+| **Post logout URI** | `https://your-bifrost-domain.com` |
+
+3. After creating the app, note the **Client ID** from the application detail page.
+
+
+ 
+
+
+4. If you chose a confidential authentication method, also copy the **Client Secret** (shown once).
+
+---
+
+## Step 2: Enable role claims on the project (Optional)
+
+
+
+
+4. Note the **Project ID** — you'll need it for the Bifrost config so Bifrost can resolve the correct project roles.
+
+
+
+
+
+1. In the same project, open the **Roles** tab and create a role for each Bifrost role you plan to map. Common pattern:
+2. Authorize users to the relevant roles via **Users → Roles**.
+
+
+ 
+
+
+
+---
+
+## Step 4: Create a service account for provisioning
+
+Web apps in Zitadel cannot use the `client_credentials` grant. Bifrost needs a dedicated service account to list users via the Management API.
+
+1. Navigate to **Users → Service Accounts → New**.
+
+
+ 
+
+
+2. Name it (e.g. `bifrost-provisioning`) and create it.
+3. Open the service account → **Actions → Generate Client Secret**.
+
+
+ 
+
+
+4. **Copy the Client ID and Client Secret immediately** — the secret is shown only once.
+
+
+
+
+---
+
+## Step 5: App token settings
+
+1. Change **Auth Token Type** to JWT.
+2. Enable roles and profile info in ID token.
+
+
+ 
+
+
+
+## Step 6: Configure Bifrost
+
+### Using the Bifrost dashboard
+
+
+1. In Bifrost, go to **Governance → User Provisioning**.
+2. Select **Zitadel** as the SCIM Provider.
+3. Fill in the fields:
+
+| Field | Value |
+| --- | --- |
+| **Domain** | Your Zitadel host, e.g. `my-instance.zitadel.cloud` or `auth.company.com` (no scheme, no path) |
+| **Project ID** | The project ID from Step 2 |
+| **Client ID** | Web Application Client ID from Step 1 |
+| **Client Secret** | Web Application Client Secret from Step 1 (optional for PKCE) |
+| **Audience** | Optional access-token audience override |
+| **Service Account Client ID** | From Step 4 |
+| **Service Account Client Secret** | From Step 4 |
+
+4. Click **Verify** — Bifrost connects to Zitadel's JWKS and service account token endpoints to confirm the credentials.
+5. Configure **Attribute → Role / Team / Business Unit** mappings if you need to translate project roles or metadata into Bifrost roles.
+6. Toggle **Enabled** and click **Save Configuration**.
+
+
+ 
+
+
+### Using `config.json`
+
+```json
+{
+ "scim_config": {
+ "enabled": true,
+ "provider": "zitadel",
+ "config": {
+ "domain": "my-instance.zitadel.cloud",
+ "projectId": "123456789012345678",
+ "clientId": "123456789012345678@my-project",
+ "clientSecret": "${ZITADEL_CLIENT_SECRET}",
+ "serviceAccountClientId": "987654321098765432@my-project",
+ "serviceAccountClientSecret": "${ZITADEL_SA_CLIENT_SECRET}",
+ "teamIdsField": "groups"
+ }
+ }
+}
+```
+
+### Custom attribute mapping
+
+You can also map any custom attributes to any entity (role, team or business unit). Make sure these are configured to send back to Bifrost in token configuration.
+
+
+
+
+
+### Configuration reference
+
+| Field | Required | Description |
+| --- | --- | --- |
+| `domain` | Yes | Zitadel instance host (no scheme). Examples: `my-instance.zitadel.cloud`, `auth.company.com`. |
+| `clientId` | Yes | Client ID of the Web Application used for user login. |
+| `clientSecret` | Yes | Web Application secret. Omit for PKCE-only flows. |
+| `projectId` | Yes | Required to resolve project-scoped role claims and sync role grants. |
+| `audience` | No | Override the expected JWT `aud` claim. |
+| `serviceAccountClientId` | Yes | Service account used to list users via the Management API. |
+| `serviceAccountClientSecret` | Yes | Service account secret (shown once in Zitadel). |
+| `attributeRoleMappings` | Yes | Ordered list of attribute→role mappings. |
+| `attributeTeamMappings` | No | Attribute→team mappings (all matches apply). |
+| `attributeBusinessUnitMappings` | No | Attribute→business-unit mappings (all matches apply). |
+
+---
+
+## Testing the Integration
+
+1. Open the Bifrost dashboard in an incognito window.
+2. You'll be redirected to Zitadel. Sign in with a user who has a project authorization.
+3. On successful login you return to Bifrost and appear in **Governance → Users** with the correct role.
+4. From **Governance → User Provisioning → Import Users**, verify you can preview and import additional users via the service account.
+
+---
+
+## Troubleshooting
+
+### `role claims missing in token`
+
+- Enable **Assert Roles on Authentication** on the project (Step 2).
+- Ensure the user has an active authorization for the project.
+
+### `invalid audience` when validating the JWT
+
+- Check the `audience` field in the Bifrost config. It must match the `aud` claim issued by Zitadel. Leaving it empty uses the default (the project's resource owner).
+
+### Service account cannot list users
+
+- Confirm the service account has **IAM_USER_READ** or **Org Owner** role in the organization.
+- Regenerate the client secret if you've lost it — the original secret cannot be retrieved.
+
+### Redirect URI mismatch
+
+- Zitadel requires an exact string match. Check for trailing slashes and `http` vs `https`.
+
+---
+
+## Next Steps
+
+- [User Provisioning overview](./user-provisioning) — capabilities, attribute mappings, bulk import
+- [Role-Based Access Control](./rbac) — configure custom roles before mapping
+- [Audit Logs](./audit-logs) — track authentication events
diff --git a/docs/enterprise/user-provisioning.mdx b/docs/enterprise/user-provisioning.mdx
new file mode 100644
index 0000000000..14a7e182b8
--- /dev/null
+++ b/docs/enterprise/user-provisioning.mdx
@@ -0,0 +1,197 @@
+---
+title: "User Provisioning (SCIM)"
+description: "Authenticate users, sync teams, and provision roles and business units from your identity provider using SCIM-backed OAuth 2.0 / OIDC flows."
+icon: "users-gear"
+---
+
+## Overview
+
+Bifrost Enterprise uses **SCIM-backed identity provisioning** to connect your organization's identity provider to Bifrost. A single configuration gives you:
+
+- **Single sign-on (SSO)** via OAuth 2.0 / OIDC with JWKS-based JWT validation
+- **Automatic role assignment** using custom claims, app roles, or group-to-role mappings
+- **Team synchronization** from IdP groups into Bifrost teams
+- **Business unit mapping** from IdP attributes to Bifrost business units
+- **Bulk user provisioning** with filter-preview before import
+- **Silent token refresh** using server-stored refresh tokens
+
+Once configured, users sign in to Bifrost with their corporate credentials and inherit the right [role and permissions](./rbac) immediately — no manual account creation.
+
+
+ 
+
+
+---
+
+## Supported Identity Providers
+
+Pick your IdP to follow a step-by-step setup guide. All providers share the same Bifrost configuration surface — the only difference is how the OAuth client and role/group claims are created on the provider side.
+
+
+
+
+1. **Login** — Bifrost redirects unauthenticated users to the provider's authorization endpoint (Authorization Code flow).
+2. **Token exchange** — on callback, Bifrost exchanges the code for an access token and refresh token, stores them in an `HttpOnly` cookie / server session, and validates the JWT against the provider's JWKS.
+3. **Identity extraction** — configurable JWT claims (`userIdField`, `rolesField`, `teamIdsField`) are mapped to a Bifrost user, role, and teams. Provider-specific app roles or custom attributes override claim lookup.
+4. **Attribute mapping** — optional `attributeRoleMappings`, `attributeTeamMappings`, and `attributeBusinessUnitMappings` translate arbitrary claim values (e.g., a department string or Okta group name) into Bifrost roles, teams, or business units.
+5. **Bulk import** — admins can preview users matching a filter and bulk-import them via the dashboard, which calls the provider's user directory API.
+6. **Silent refresh** — when the access token expires, Bifrost uses the stored refresh token to mint a new one without requiring re-login.
+
+---
+
+## Capabilities
+
+| Capability | Description |
+| --- | --- |
+| **OAuth 2.0 / OIDC SSO** | Authorization Code + PKCE with configurable scopes (`openid profile email offline_access`). |
+| **JWKS validation** | JWTs are validated against the provider's published JWKS keys; configuration is cached and auto-refreshed. |
+| **Role mapping** | Map from a claim value (string or array) to Admin / Developer / Viewer or a custom role. Highest-privilege wins when multiple match. |
+| **Team mapping** | Map multiple claim values to Bifrost teams in a single pass (a user can belong to many teams). |
+| **Business unit mapping** | Same as team mapping but scoped to business units. |
+| **Provisioning preview** | Preview up to 50 users matching filters (groups, roles, departments) before importing. |
+| **Bulk import** | Import matched users into Bifrost with role + team + BU assignments applied. |
+| **Team sync** | Sync IdP groups as Bifrost teams with a single action. |
+| **Business unit sync** | Sync IdP organizational units as Bifrost business units. |
+| **Deprovisioning** | Re-running import reconciles removed users and updates role / team assignments. |
+| **API key pass-through** | Requests using Bifrost API keys (`bfst-*`) bypass SCIM middleware so inference traffic is not affected. |
+
+---
+
+## Configuration reference
+
+All providers share the same outer config shape in `config.json`:
+
+```json
+{
+ "scim_config": {
+ "enabled": true,
+ "provider": "okta | entra | zitadel | keycloak | google",
+ "config": {
+ "...": "provider-specific fields — see each IdP guide"
+ }
+ }
+}
+```
+
+Shared fields across providers:
+
+| Field | Required | Description |
+| --- | --- | --- |
+| `clientId` | Yes | OAuth client ID from the identity provider. |
+| `clientSecret` | Usually | Client secret. Required for confidential clients and (where applicable) token revocation. |
+| `audience` | Optional | JWT audience to validate against. Defaults vary per provider. |
+| `attributeRoleMappings` | Optional | Ordered list of `{ attribute, value, role }` rules evaluated top-to-bottom. |
+| `attributeTeamMappings` | Optional | List of `{ attribute, value, team }` rules (all matches apply). |
+| `attributeBusinessUnitMappings` | Optional | List of `{ attribute, value, businessUnit }` rules (all matches apply). |
+
+Provider-specific fields (domain, tenant ID, server URL, service-account credentials) are documented in each IdP's setup guide.
+
+
+
+
+4. Click **Verify** to test credentials end-to-end. Bifrost will reach the provider's JWKS / directory endpoint and report any failures.
+5. Configure **Attribute → Role / Team / Business Unit** mappings as needed.
+6. Toggle **Enabled** and click **Save Configuration**.
+
+
+
+
+Each mapping is an ordered rule:
+
+```json
+{
+ "attribute": "department",
+ "value": "Engineering",
+ "role": "developer"
+}
+```
+
+Rules are evaluated top-to-bottom:
+- **Role mappings** — first match wins. Set a fallback with `"attribute": "*"` at the end.
+- **Team mappings** and **business unit mappings** — all matching rules apply, so a user with `department=Platform` and `group=sre` can be placed on multiple teams.
+
+Claim values can be strings, arrays, or nested objects — Bifrost resolves dotted paths (e.g., `realm_access.roles`).
+
+---
+
+## Bulk user provisioning
+
+Once SCIM is enabled, import users in bulk from your IdP:
+
+1. Go to **Governance → User Provisioning → Import Users**.
+2. Select a filter — groups, roles, departments, or a custom query depending on provider support.
+3. Click **Preview** to see up to 50 matching users.
+4. Click **Import** to create them in Bifrost with role / team / BU assignments applied.
+
+
+ 
+
+
+Re-running an import reconciles existing users — role and team changes in the IdP are reflected on the next import.
+
+---
+
+## Troubleshooting
+
+| Symptom | Likely cause |
+| --- | --- |
+| Access denied: no application role or group mapping is assigned to this user. | Make sure you have assigned user to the Bifrost IdP application and they have a valid group/attribute mapping to role in Bifrost |
+| Redirect loop on login | Make sure you have restarted pods/Bifrost instance after changing SCIM configuration, or check for a redirect URI mismatch. Exact string match required — check trailing slashes and `http` vs `https`. |
+| `invalid audience` | `audience` field does not match the access token's `aud` claim. Use the same value your IdP issues. |
+| Empty roles / teams | Claim mapping is off. Verify the JWT at [jwt.io](https://jwt.io) and check `rolesField` / `teamIdsField`. |
+| Token refresh failing | `offline_access` scope missing or refresh token revoked. Re-enable the scope and re-authenticate. |
+| First user gets Admin | By design — if no matching role mapping applies, the first user is promoted to Admin so they can finish configuration. Subsequent users default to Viewer. |
+
+Provider-specific troubleshooting lives in each IdP's guide.
+
+---
+
+## Related
+
+- [Role-Based Access Control](./rbac) — permissions model and custom roles
+- [Advanced Governance](./advanced-governance) — budgets, limits, and compliance
+- [Audit Logs](./audit-logs) — track authentication events and role changes
diff --git a/docs/features/governance/budget-and-limits.mdx b/docs/features/governance/budget-and-limits.mdx
index 1a012b7904..b295565e72 100644
--- a/docs/features/governance/budget-and-limits.mdx
+++ b/docs/features/governance/budget-and-limits.mdx
@@ -345,59 +345,78 @@ Configure provider-level governance through Bifrost's configuration file for dec
"governance": {
"virtual_keys": [
{
+ "id": "vk-dev-001",
"name": "development-team-vk",
"description": "Development team with multi-provider setup",
+ "is_active": true,
+ "rate_limit_id": "rl-vk-dev",
"provider_configs": [
{
+ "id": 1,
"provider": "openai",
"weight": 0.6,
"allowed_models": ["gpt-4", "gpt-3.5-turbo"],
- "budget": {
- "max_limit": 1000.00,
- "reset_duration": "1M"
- },
- "rate_limit": {
- "token_max_limit": 2000000,
- "token_reset_duration": "1h",
- "request_max_limit": 2000,
- "request_reset_duration": "1h"
- }
+ "rate_limit_id": "rl-pc-openai"
},
{
+ "id": 2,
"provider": "anthropic",
"weight": 0.4,
"allowed_models": ["claude-3-opus", "claude-3-sonnet"],
- "budget": {
- "max_limit": 500.00,
- "reset_duration": "1M"
- },
- "rate_limit": {
- "token_max_limit": 1000000,
- "token_reset_duration": "1h",
- "request_max_limit": 1000,
- "request_reset_duration": "1h"
- }
+ "rate_limit_id": "rl-pc-anthropic"
}
- ],
- "budget": {
- "max_limit": 2000.00,
- "reset_duration": "1M",
- "calendar_aligned": true
- },
- "rate_limit": {
- "token_max_limit": 5000000,
- "token_reset_duration": "1h",
- "request_max_limit": 3000,
- "request_reset_duration": "1h"
- },
- "is_active": true
+ ]
+ }
+ ],
+ "budgets": [
+ {
+ "id": "budget-vk-dev",
+ "virtual_key_id": "vk-dev-001",
+ "max_limit": 2000.00,
+ "reset_duration": "1M",
+ "calendar_aligned": true
+ },
+ {
+ "id": "budget-pc-openai",
+ "provider_config_id": 1,
+ "max_limit": 1000.00,
+ "reset_duration": "1M"
+ },
+ {
+ "id": "budget-pc-anthropic",
+ "provider_config_id": 2,
+ "max_limit": 500.00,
+ "reset_duration": "1M"
+ }
+ ],
+ "rate_limits": [
+ {
+ "id": "rl-vk-dev",
+ "token_max_limit": 5000000,
+ "token_reset_duration": "1h",
+ "request_max_limit": 3000,
+ "request_reset_duration": "1h"
+ },
+ {
+ "id": "rl-pc-openai",
+ "token_max_limit": 2000000,
+ "token_reset_duration": "1h",
+ "request_max_limit": 2000,
+ "request_reset_duration": "1h"
+ },
+ {
+ "id": "rl-pc-anthropic",
+ "token_max_limit": 1000000,
+ "token_reset_duration": "1h",
+ "request_max_limit": 1000,
+ "request_reset_duration": "1h"
}
]
}
}
```
-Optional `calendar_aligned` on each `budget` matches the HTTP API and [calendar-aligned behavior](#calendar-aligned-budgets).
+Budgets and rate limits live as **separate top-level arrays** inside `governance`. Virtual keys and provider configs reference them by id (`rate_limit_id`) or are referenced back (`virtual_key_id` / `provider_config_id` on each `budgets[]` entry). Optional `calendar_aligned` on each `budget` matches the HTTP API and [calendar-aligned behavior](#calendar-aligned-budgets).
### Advanced Configuration Examples
@@ -407,34 +426,21 @@ Optional `calendar_aligned` on each `budget` matches the HTTP API and [calendar-
"governance": {
"virtual_keys": [
{
+ "id": "vk-cost-opt",
"name": "cost-optimized-vk",
"provider_configs": [
- {
- "provider": "openai-gpt-3.5",
- "weight": 0.8,
- "budget": {
- "max_limit": 50.00,
- "reset_duration": "1d"
- },
- "rate_limit": {
- "request_max_limit": 1000,
- "request_reset_duration": "1h"
- }
- },
- {
- "provider": "openai-gpt-4",
- "weight": 0.2,
- "budget": {
- "max_limit": 200.00,
- "reset_duration": "1d"
- },
- "rate_limit": {
- "request_max_limit": 100,
- "request_reset_duration": "1h"
- }
- }
+ {"id": 10, "provider": "openai-gpt-3.5", "weight": 0.8, "rate_limit_id": "rl-cheap"},
+ {"id": 11, "provider": "openai-gpt-4", "weight": 0.2, "rate_limit_id": "rl-premium"}
]
}
+ ],
+ "budgets": [
+ {"id": "b-cheap", "provider_config_id": 10, "max_limit": 50.00, "reset_duration": "1d"},
+ {"id": "b-premium", "provider_config_id": 11, "max_limit": 200.00, "reset_duration": "1d"}
+ ],
+ "rate_limits": [
+ {"id": "rl-cheap", "request_max_limit": 1000, "request_reset_duration": "1h"},
+ {"id": "rl-premium", "request_max_limit": 100, "request_reset_duration": "1h"}
]
}
}
@@ -446,52 +452,24 @@ Optional `calendar_aligned` on each `budget` matches the HTTP API and [calendar-
"governance": {
"virtual_keys": [
{
+ "id": "vk-prod-hv",
"name": "production-high-volume-vk",
"provider_configs": [
- {
- "provider": "openai",
- "weight": 0.5,
- "budget": {
- "max_limit": 5000.00,
- "reset_duration": "1M"
- },
- "rate_limit": {
- "token_max_limit": 10000000,
- "token_reset_duration": "1h",
- "request_max_limit": 10000,
- "request_reset_duration": "1h"
- }
- },
- {
- "provider": "anthropic",
- "weight": 0.3,
- "budget": {
- "max_limit": 3000.00,
- "reset_duration": "1M"
- },
- "rate_limit": {
- "token_max_limit": 6000000,
- "token_reset_duration": "1h",
- "request_max_limit": 6000,
- "request_reset_duration": "1h"
- }
- },
- {
- "provider": "azure-openai",
- "weight": 0.2,
- "budget": {
- "max_limit": 2000.00,
- "reset_duration": "1M"
- },
- "rate_limit": {
- "token_max_limit": 4000000,
- "token_reset_duration": "1h",
- "request_max_limit": 4000,
- "request_reset_duration": "1h"
- }
- }
+ {"id": 20, "provider": "openai", "weight": 0.5, "rate_limit_id": "rl-openai"},
+ {"id": 21, "provider": "anthropic", "weight": 0.3, "rate_limit_id": "rl-anthropic"},
+ {"id": 22, "provider": "azure-openai", "weight": 0.2, "rate_limit_id": "rl-azure"}
]
}
+ ],
+ "budgets": [
+ {"id": "b-openai", "provider_config_id": 20, "max_limit": 5000.00, "reset_duration": "1M"},
+ {"id": "b-anthropic", "provider_config_id": 21, "max_limit": 3000.00, "reset_duration": "1M"},
+ {"id": "b-azure", "provider_config_id": 22, "max_limit": 2000.00, "reset_duration": "1M"}
+ ],
+ "rate_limits": [
+ {"id": "rl-openai", "token_max_limit": 10000000, "token_reset_duration": "1h", "request_max_limit": 10000, "request_reset_duration": "1h"},
+ {"id": "rl-anthropic", "token_max_limit": 6000000, "token_reset_duration": "1h", "request_max_limit": 6000, "request_reset_duration": "1h"},
+ {"id": "rl-azure", "token_max_limit": 4000000, "token_reset_duration": "1h", "request_max_limit": 4000, "request_reset_duration": "1h"}
]
}
}
@@ -514,20 +492,23 @@ A virtual key configured with multiple providers and different budget allocation
```json
{
- "name": "marketing-team-vk",
- "budget": { "max_limit": 100, "reset_duration": "1M" },
- "provider_configs": [
- {
- "provider": "openai",
- "weight": 0.7,
- "budget": { "max_limit": 50, "reset_duration": "1M" }
- },
- {
- "provider": "anthropic",
- "weight": 0.3,
- "budget": { "max_limit": 30, "reset_duration": "1M" }
- }
- ]
+ "governance": {
+ "virtual_keys": [
+ {
+ "id": "vk-mkt",
+ "name": "marketing-team-vk",
+ "provider_configs": [
+ {"id": 30, "provider": "openai", "weight": 0.7},
+ {"id": 31, "provider": "anthropic", "weight": 0.3}
+ ]
+ }
+ ],
+ "budgets": [
+ {"id": "b-vk-mkt", "virtual_key_id": "vk-mkt", "max_limit": 100, "reset_duration": "1M"},
+ {"id": "b-openai", "provider_config_id": 30, "max_limit": 50, "reset_duration": "1M"},
+ {"id": "b-anth", "provider_config_id": 31, "max_limit": 30, "reset_duration": "1M"}
+ ]
+ }
}
```
@@ -542,27 +523,22 @@ Different rate limits based on provider capabilities:
```json
{
- "name": "high-volume-vk",
- "provider_configs": [
- {
- "provider": "openai",
- "rate_limit": {
- "request_max_limit": 1000,
- "request_reset_duration": "1h",
- "token_max_limit": 1000000,
- "token_reset_duration": "1h"
- }
- },
- {
- "provider": "anthropic",
- "rate_limit": {
- "request_max_limit": 500,
- "request_reset_duration": "1h",
- "token_max_limit": 500000,
- "token_reset_duration": "1h"
+ "governance": {
+ "virtual_keys": [
+ {
+ "id": "vk-hv",
+ "name": "high-volume-vk",
+ "provider_configs": [
+ {"id": 40, "provider": "openai", "rate_limit_id": "rl-openai"},
+ {"id": 41, "provider": "anthropic", "rate_limit_id": "rl-anthropic"}
+ ]
}
- }
- ]
+ ],
+ "rate_limits": [
+ {"id": "rl-openai", "request_max_limit": 1000, "request_reset_duration": "1h", "token_max_limit": 1000000, "token_reset_duration": "1h"},
+ {"id": "rl-anthropic", "request_max_limit": 500, "request_reset_duration": "1h", "token_max_limit": 500000, "token_reset_duration": "1h"}
+ ]
+ }
}
```
@@ -577,25 +553,25 @@ Provider configurations with budget-based failover:
```json
{
- "name": "cost-optimized-vk",
- "provider_configs": [
- {
- "provider": "openai-cheap",
- "weight": 1.0,
- "budget": { "max_limit": 10, "reset_duration": "1d" }
- },
- {
- "provider": "openai-premium",
- "weight": 0.0,
- "budget": { "max_limit": 50, "reset_duration": "1d" },
- "rate_limit": {
- "request_max_limit": 100,
- "request_reset_duration": "1h",
- "token_max_limit": 50000,
- "token_reset_duration": "1h"
+ "governance": {
+ "virtual_keys": [
+ {
+ "id": "vk-cost",
+ "name": "cost-optimized-vk",
+ "provider_configs": [
+ {"id": 50, "provider": "openai-cheap", "weight": 1.0},
+ {"id": 51, "provider": "openai-premium", "weight": 0.0, "rate_limit_id": "rl-premium"}
+ ]
}
- }
- ]
+ ],
+ "budgets": [
+ {"id": "b-cheap", "provider_config_id": 50, "max_limit": 10, "reset_duration": "1d"},
+ {"id": "b-premium", "provider_config_id": 51, "max_limit": 50, "reset_duration": "1d"}
+ ],
+ "rate_limits": [
+ {"id": "rl-premium", "request_max_limit": 100, "request_reset_duration": "1h", "token_max_limit": 50000, "token_reset_duration": "1h"}
+ ]
+ }
}
```
diff --git a/docs/features/governance/virtual-keys.mdx b/docs/features/governance/virtual-keys.mdx
index 8ae28e42ce..6c8cff9320 100644
--- a/docs/features/governance/virtual-keys.mdx
+++ b/docs/features/governance/virtual-keys.mdx
@@ -169,7 +169,8 @@ curl -X DELETE http://localhost:8080/api/governance/virtual-keys/{vk_id}
{
"provider": "openai",
"weight": 0.5,
- "allowed_models": ["gpt-4o-mini"]
+ "allowed_models": ["gpt-4o-mini"],
+ "key_ids": ["openai-primary"]
},
{
"provider": "anthropic",
@@ -178,11 +179,7 @@ curl -X DELETE http://localhost:8080/api/governance/virtual-keys/{vk_id}
}
],
"team_id": "team-eng-001",
- "budget_id": "budget-eng-vk",
- "rate_limit_id": "rate-limit-eng-vk",
- "keys": [
- {"key_id": "8c52039e-38c6-48b2-8016-0bd884b7befb"}
- ]
+ "rate_limit_id": "rate-limit-eng-vk"
},
{
"id": "vk-002",
@@ -202,16 +199,13 @@ curl -X DELETE http://localhost:8080/api/governance/virtual-keys/{vk_id}
"allowed_models": ["claude-3-opus-20240229"]
}
],
- "customer_id": "customer-acme-corp",
- "budget_id": "budget-exec-vk",
- "keys": [
- {"key_id": "8c52039e-38c6-48b2-8016-0bd884b7befb"}
- ]
+ "customer_id": "customer-acme-corp"
}
],
"budgets": [
{
"id": "budget-eng-vk",
+ "virtual_key_id": "vk-001",
"max_limit": 100.00,
"reset_duration": "1M",
"current_usage": 0.0,
@@ -219,6 +213,7 @@ curl -X DELETE http://localhost:8080/api/governance/virtual-keys/{vk_id}
},
{
"id": "budget-exec-vk",
+ "virtual_key_id": "vk-002",
"max_limit": 500.00,
"reset_duration": "1M",
"current_usage": 0.0,
diff --git a/docs/features/litellm-compat.mdx b/docs/features/litellm-compat.mdx
index 51cd26dcd9..490a37efa4 100644
--- a/docs/features/litellm-compat.mdx
+++ b/docs/features/litellm-compat.mdx
@@ -9,8 +9,10 @@ icon: "train"
The LiteLLM compatibility plugin provides two transformations:
1. **Text-to-Chat Conversion** - Automatically converts text completion requests to chat completion format for models that only support chat APIs
+2. **Chat-to-Responses Conversion** - Automatically converts chat completion requests to responses format for models that only support responses APIs
+3. **Drop Unsupported Params** - Automatically drops unsupported parameters if the model doesn't support them
-When either transformation is applied, responses include `extra_fields.litellm_compat: true`.
+When either transformation is applied, responses include `extra_fields.converted_request_type: 
+
+
+- Copy API key to use it in the Azure content moderation config form
+- Copy project endpoint and use base URL as endpoint in the form. e.g. (`https://xxx-resource.services.ai.azure.com`)
## Severity Threshold Levels
-| Threshold | Numeric Value | Behavior |
-|-----------|---------------|----------|
-| `low` | 2 | Most strict - blocks severity 2 and above |
-| `medium` | 4 | Balanced - blocks severity 4 and above |
-| `high` | 6 | Least strict - blocks only severity 6 |
+| Threshold | Numeric Value | Behavior |
+| --------- | ------------- | ----------------------------------------- |
+| `low` | 2 | Most strict - blocks severity 2 and above |
+| `medium` | 4 | Balanced - blocks severity 4 and above |
+| `high` | 6 | Least strict - blocks only severity 6 |
## Detection Categories
@@ -47,23 +58,23 @@ Bifrost integrates with **Azure AI Content Safety** to provide multi-modal conte
- Self-harm