diff --git a/core/bifrost.go b/core/bifrost.go index 259b0cabea9..97e6df7d4ac 100644 --- a/core/bifrost.go +++ b/core/bifrost.go @@ -5184,6 +5184,19 @@ func (bifrost *Bifrost) handleStreamRequest(ctx *schemas.BifrostContext, req *sc // the channel. See SetFallbackRoutingInfo doc. fallbackErr.SetFallbackRoutingInfo(provider, model) if fallbackErr == nil { + // Layer Primary/IsFallback onto the ctx-stashed RoutingInfo (mirrors + // SetFallbackRoutingInfo) — chunks can't carry it (see above), but the + // transport's pre-stream response headers can. + if ri, ok := ctx.Value(schemas.BifrostContextKeyRoutingInfo).(schemas.RoutingInfo); ok { + ri.IsFallback = true + if provider != "" { + ri.PrimaryProvider = new(provider) + } + if model != "" { + ri.PrimaryModel = new(model) + } + ctx.SetValue(schemas.BifrostContextKeyRoutingInfo, ri) + } bifrost.logger.Debug(fmt.Sprintf("successfully used fallback provider %s with model %s", fallback.Provider, fallback.Model)) ctx.AppendRoutingEngineLog(schemas.RoutingEngineCore, schemas.LogLevelInfo, fmt.Sprintf("Request served by fallback %s/%s (attempt %d/%d)", fallback.Provider, fallback.Model, i+1, len(fallbacks))) tracer.EndSpan(handle, schemas.SpanStatusOk, "") @@ -6674,6 +6687,10 @@ func (bifrost *Bifrost) requestWorker(provider schemas.Provider, config *schemas // alias while this attempt's provider goroutine is still emitting chunks. attemptResolvedModel := resolvedModel attemptRoutingInfo = schemas.BuildRoutingInfo(req.Context, provider.GetProviderKey(), originalModelRequested, k) + // Stash for the transport: streams carry RoutingInfo only on chunks, but + // response headers must be written before the first chunk arrives. Each + // retry overwrites, so the winning attempt's snapshot survives. + req.Context.SetValue(schemas.BifrostContextKeyRoutingInfo, attemptRoutingInfo) // Per-attempt snapshot for the async postHookRunner closure (it must // not capture the outer var by reference — a later retry would race). perAttemptRoutingInfo := attemptRoutingInfo diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 81c8a5d4411..b4a717bdcde 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -261,6 +261,7 @@ const ( BifrostContextKeyNumberOfRetries BifrostContextKey = "bifrost-number-of-retries" // int (to store the number of retries (set by bifrost - DO NOT SET THIS MANUALLY)) BifrostContextKeyFallbackIndex BifrostContextKey = "bifrost-fallback-index" // int (to store the fallback index (set by bifrost - DO NOT SET THIS MANUALLY)) 0 for primary, 1 for first fallback, etc. BifrostContextKeyResolvedAlias BifrostContextKey = "bifrost-resolved-alias" // *ResolvedAlias (set by bifrost after key-level alias resolution — providers read this for model_family routing and provider-specific overrides; nil/absent when no alias matched) + BifrostContextKeyRoutingInfo BifrostContextKey = "bifrost-routing-info" // RoutingInfo (set by bifrost per stream attempt - DO NOT SET THIS MANUALLY) - streams carry RoutingInfo only on chunks, so the transport reads this snapshot to emit routed-identity response headers before the first chunk BifrostContextKeyStreamEndIndicator BifrostContextKey = "bifrost-stream-end-indicator" // bool (set by bifrost - DO NOT SET THIS MANUALLY) BifrostContextKeyStreamGated BifrostContextKey = "bifrost-stream-gated" // bool (set by ctx.PauseStream/ResumeStream/EndStream when a plugin first engages the pause/resume gate; provider helpers use this as a fast-path check to skip Tracer.GateSend on streams that never engage the gate) BifrostContextKeyStreamIdleTimeout BifrostContextKey = "bifrost-stream-idle-timeout" // time.Duration (per-chunk idle timeout for streaming) @@ -1257,6 +1258,19 @@ func (e *BifrostError) PopulateRoutingInfo(info RoutingInfo) { syncDeprecatedFromRoutingInfo(info, &e.ExtraFields.Provider, &e.ExtraFields.OriginalModelRequested, &e.ExtraFields.ResolvedModelUsed) } +// ToExtraFields builds a response-shaped ExtraFields snapshot from a finalized +// RoutingInfo, deriving the deprecated triplet via the same rules as the +// response path. Used by the transport to emit routed-identity headers for +// streams, where per-chunk ExtraFields don't exist yet at header-write time. +func (ri RoutingInfo) ToExtraFields(requestType RequestType) BifrostResponseExtraFields { + extra := BifrostResponseExtraFields{ + RequestType: requestType, + RoutingInfo: ri, + } + syncDeprecatedFromRoutingInfo(ri, &extra.Provider, &extra.OriginalModelRequested, &extra.ResolvedModelUsed) + return extra +} + // SetFallbackRoutingInfo marks the active sub-response's RoutingInfo as a // fallback attempt and records the primary attempt's provider/model. Also // re-syncs the deprecated OriginalModelRequested to the primary model per diff --git a/core/schemas/context.go b/core/schemas/context.go index 8ed6d037467..414a3553685 100644 --- a/core/schemas/context.go +++ b/core/schemas/context.go @@ -30,6 +30,7 @@ var reservedKeys = []any{ BifrostContextKeyStreamGated, BifrostContextKeyMCPHealthCheckRequest, BifrostContextKeyUpstreamLatency, + BifrostContextKeyRoutingInfo, } // pluginLogStore holds plugin log entries accumulated during request processing. diff --git a/transports/bifrost-http/handlers/inference.go b/transports/bifrost-http/handlers/inference.go index 32eb1121a25..b0bf50aae71 100644 --- a/transports/bifrost-http/handlers/inference.go +++ b/transports/bifrost-http/handlers/inference.go @@ -1862,7 +1862,7 @@ func (h *CompletionHandler) handleStreamingTextCompletion(ctx *fasthttp.RequestC return h.client.TextCompletionStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.TextCompletionStreamRequest, getStream, cancel) } // handleStreamingChatCompletion handles streaming chat completion requests using Server-Sent Events (SSE) @@ -1874,7 +1874,7 @@ func (h *CompletionHandler) handleStreamingChatCompletion(ctx *fasthttp.RequestC return h.client.ChatCompletionStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.ChatCompletionStreamRequest, getStream, cancel) } // handleStreamingResponses handles streaming responses requests using Server-Sent Events (SSE) @@ -1886,7 +1886,7 @@ func (h *CompletionHandler) handleStreamingResponses(ctx *fasthttp.RequestCtx, r return h.client.ResponsesStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.ResponsesStreamRequest, getStream, cancel) } // handleStreamingSpeech handles streaming speech requests using Server-Sent Events (SSE) @@ -1898,7 +1898,7 @@ func (h *CompletionHandler) handleStreamingSpeech(ctx *fasthttp.RequestCtx, req return h.client.SpeechStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.SpeechStreamRequest, getStream, cancel) } // handleStreamingTranscriptionRequest handles streaming transcription requests using Server-Sent Events (SSE) @@ -1910,14 +1910,14 @@ func (h *CompletionHandler) handleStreamingTranscriptionRequest(ctx *fasthttp.Re return h.client.TranscriptionStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.TranscriptionStreamRequest, getStream, cancel) } // handleStreamingResponse is a generic function to handle streaming responses using Server-Sent Events (SSE) // The cancel function is called ONLY when client disconnects are detected via write errors. // Bifrost handles cleanup internally for normal completion and errors, so we only cancel // upstream streams when write errors indicate the client has disconnected. -func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, getStream func() (chan *schemas.BifrostStreamChunk, *schemas.BifrostError), cancel context.CancelFunc) { +func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, requestType schemas.RequestType, getStream func() (chan *schemas.BifrostStreamChunk, *schemas.BifrostError), cancel context.CancelFunc) { // Get the streaming channel — called BEFORE setting SSE headers so that // provider errors return proper HTTP status codes + JSON content type. stream, bifrostErr := getStream() @@ -1938,6 +1938,10 @@ func (h *CompletionHandler) handleStreamingResponse(ctx *fasthttp.RequestCtx, bi forwardProviderHeaders(ctx, headers) } + // Routed-identity headers from the context snapshot — routing is final once + // the stream channel is returned, before any chunk arrives. + lib.ApplyBifrostStreamResponseHeaders(ctx, bifrostCtx, requestType) + // Signal to tracing middleware that trace completion should be deferred // The streaming callback will complete the trace after the stream ends ctx.SetUserValue(schemas.BifrostContextKeyDeferTraceCompletion, true) @@ -2286,7 +2290,7 @@ func (h *CompletionHandler) handleStreamingImageGeneration(ctx *fasthttp.Request return h.client.ImageGenerationStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.ImageGenerationStreamRequest, getStream, cancel) } // prepareImageEditRequest prepares a BifrostImageEditRequest from a multipart form @@ -2492,7 +2496,7 @@ func (h *CompletionHandler) handleStreamingImageEditRequest(ctx *fasthttp.Reques return h.client.ImageEditStreamRequest(bifrostCtx, req) } - h.handleStreamingResponse(ctx, bifrostCtx, getStream, cancel) + h.handleStreamingResponse(ctx, bifrostCtx, schemas.ImageEditStreamRequest, getStream, cancel) } // prepareImageVariationRequest prepares a BifrostImageVariationRequest from a multipart form diff --git a/transports/bifrost-http/integrations/router.go b/transports/bifrost-http/integrations/router.go index 8493e7dda9c..01acbfb5009 100644 --- a/transports/bifrost-http/integrations/router.go +++ b/transports/bifrost-http/integrations/router.go @@ -2564,21 +2564,29 @@ func (g *GenericRouter) handleStreamingRequest(ctx *fasthttp.RequestCtx, config // We now get a cancellable context from ConvertToBifrostContext so we can cancel the upstream stream immediately when the client disconnects. var stream chan *schemas.BifrostStreamChunk var bifrostErr *schemas.BifrostError + var requestType schemas.RequestType // Handle different request types if bifrostReq.TextCompletionRequest != nil { + requestType = schemas.TextCompletionStreamRequest stream, bifrostErr = g.client.TextCompletionStreamRequest(bifrostCtx, bifrostReq.TextCompletionRequest) } else if bifrostReq.ChatRequest != nil { + requestType = schemas.ChatCompletionStreamRequest stream, bifrostErr = g.client.ChatCompletionStreamRequest(bifrostCtx, bifrostReq.ChatRequest) } else if bifrostReq.ResponsesRequest != nil { + requestType = schemas.ResponsesStreamRequest stream, bifrostErr = g.client.ResponsesStreamRequest(bifrostCtx, bifrostReq.ResponsesRequest) } else if bifrostReq.SpeechRequest != nil { + requestType = schemas.SpeechStreamRequest stream, bifrostErr = g.client.SpeechStreamRequest(bifrostCtx, bifrostReq.SpeechRequest) } else if bifrostReq.TranscriptionRequest != nil { + requestType = schemas.TranscriptionStreamRequest stream, bifrostErr = g.client.TranscriptionStreamRequest(bifrostCtx, bifrostReq.TranscriptionRequest) } else if bifrostReq.ImageGenerationRequest != nil { + requestType = schemas.ImageGenerationStreamRequest stream, bifrostErr = g.client.ImageGenerationStreamRequest(bifrostCtx, bifrostReq.ImageGenerationRequest) } else if bifrostReq.ImageEditRequest != nil { + requestType = schemas.ImageEditStreamRequest stream, bifrostErr = g.client.ImageEditStreamRequest(bifrostCtx, bifrostReq.ImageEditRequest) } @@ -2605,6 +2613,10 @@ func (g *GenericRouter) handleStreamingRequest(ctx *fasthttp.RequestCtx, config } } + // Routed-identity headers from the context snapshot — routing is final once + // the stream channel is returned, before any chunk arrives. + lib.ApplyBifrostStreamResponseHeaders(ctx, bifrostCtx, requestType) + // Large payload streaming passthrough — bypass SSE event processing, pipe raw upstream if g.tryStreamLargeResponse(ctx, bifrostCtx, schemas.BifrostResponseExtraFields{}) { ctx.Response.Header.Set("Cache-Control", "no-cache") diff --git a/transports/bifrost-http/lib/responseheaders.go b/transports/bifrost-http/lib/responseheaders.go index 74dcecd187a..b245d46bfff 100644 --- a/transports/bifrost-http/lib/responseheaders.go +++ b/transports/bifrost-http/lib/responseheaders.go @@ -55,6 +55,23 @@ const ( HeaderBifrostRoutingInfoServerSideFallbackModel = "x-bifrost-routing-info-server-side-fallback-model" ) +// ApplyBifrostStreamResponseHeaders emits the routed-identity headers for a +// streaming response, before the first SSE write. Streams only carry +// ExtraFields on chunks — none exist at header-write time — so the identity +// comes from the RoutingInfo snapshot core stashes in the context at stream +// setup (BifrostContextKeyRoutingInfo). Absent snapshot (e.g. a plugin +// short-circuited the stream) emits only the request-type header. +func ApplyBifrostStreamResponseHeaders(ctx *fasthttp.RequestCtx, bifrostCtx *schemas.BifrostContext, requestType schemas.RequestType) { + if bifrostCtx == nil { + return + } + extra := schemas.BifrostResponseExtraFields{RequestType: requestType} + if ri, ok := bifrostCtx.Value(schemas.BifrostContextKeyRoutingInfo).(schemas.RoutingInfo); ok { + extra = ri.ToExtraFields(requestType) + } + ApplyBifrostResponseHeaders(ctx, bifrostCtx, extra) +} + // ApplyBifrostResponseHeaders writes both the upstream provider response // headers (forwarded verbatim) and the bifrost-level `x-bifrost-*` routing // identity headers onto the fasthttp response. Empty fields are skipped so diff --git a/transports/bifrost-http/lib/responseheaders_test.go b/transports/bifrost-http/lib/responseheaders_test.go index a60f29b239b..9ee422ee985 100644 --- a/transports/bifrost-http/lib/responseheaders_test.go +++ b/transports/bifrost-http/lib/responseheaders_test.go @@ -174,3 +174,76 @@ func TestApplyBifrostResponseHeaders(t *testing.T) { assert.Empty(t, string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoServerSideFallbackModel))) }) } + +// TestApplyBifrostStreamResponseHeaders covers the streaming variant: identity +// comes from the RoutingInfo snapshot core stashes in the context at stream +// setup, since no chunk (and hence no ExtraFields) exists at header-write time. +func TestApplyBifrostStreamResponseHeaders(t *testing.T) { + newBifrostCtx := func() *schemas.BifrostContext { + return schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + } + + t.Run("context snapshot emits identity and derived deprecated headers", func(t *testing.T) { + ctx := &fasthttp.RequestCtx{} + bifrostCtx := newBifrostCtx() + + bifrostCtx.SetValue(schemas.BifrostContextKeyRoutingInfo, schemas.RoutingInfo{ + Provider: schemas.Bedrock, + Model: "claude-sonnet-4-6", + Key: "prod-key-1", + ResolvedKeyAlias: &schemas.ResolvedKeyAlias{ + ModelID: "us.anthropic.claude-sonnet-4-6", + }, + }) + + ApplyBifrostStreamResponseHeaders(ctx, bifrostCtx, schemas.ChatCompletionStreamRequest) + + assert.Equal(t, string(schemas.ChatCompletionStreamRequest), string(ctx.Response.Header.Peek(HeaderBifrostRequestType))) + assert.Equal(t, "bedrock", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoProvider))) + assert.Equal(t, "claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoModel))) + assert.Equal(t, "prod-key-1", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoKey))) + assert.Equal(t, "us.anthropic.claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoAliasModelID))) + // Deprecated triplet derived from RoutingInfo via the shared sync rules. + assert.Equal(t, "bedrock", string(ctx.Response.Header.Peek(HeaderBifrostProvider))) + assert.Equal(t, "claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostOriginalModel))) + assert.Equal(t, "us.anthropic.claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostResolvedModel))) + }) + + t.Run("fallback-layered snapshot emits fallback headers", func(t *testing.T) { + ctx := &fasthttp.RequestCtx{} + bifrostCtx := newBifrostCtx() + + primaryProvider := schemas.Anthropic + primaryModel := "claude-sonnet-4-6" + bifrostCtx.SetValue(schemas.BifrostContextKeyRoutingInfo, schemas.RoutingInfo{ + Provider: schemas.Bedrock, + Model: "us.anthropic.claude-sonnet-4-6", + Key: "bedrock-key", + IsFallback: true, + PrimaryProvider: &primaryProvider, + PrimaryModel: &primaryModel, + }) + bifrostCtx.SetValue(schemas.BifrostContextKeyFallbackIndex, 1) + + ApplyBifrostStreamResponseHeaders(ctx, bifrostCtx, schemas.ChatCompletionStreamRequest) + + assert.Equal(t, "true", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoIsFallback))) + assert.Equal(t, "anthropic", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoPrimaryProvider))) + assert.Equal(t, "claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoPrimaryModel))) + assert.Equal(t, "1", string(ctx.Response.Header.Peek(HeaderBifrostFallbackIndex))) + // Deprecated original-model derives from the primary on fallback. + assert.Equal(t, "claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostOriginalModel))) + assert.Equal(t, "us.anthropic.claude-sonnet-4-6", string(ctx.Response.Header.Peek(HeaderBifrostResolvedModel))) + }) + + t.Run("missing snapshot emits only request type", func(t *testing.T) { + ctx := &fasthttp.RequestCtx{} + bifrostCtx := newBifrostCtx() + + ApplyBifrostStreamResponseHeaders(ctx, bifrostCtx, schemas.ResponsesStreamRequest) + + assert.Equal(t, string(schemas.ResponsesStreamRequest), string(ctx.Response.Header.Peek(HeaderBifrostRequestType))) + assert.Empty(t, string(ctx.Response.Header.Peek(HeaderBifrostRoutingInfoProvider))) + assert.Empty(t, string(ctx.Response.Header.Peek(HeaderBifrostProvider))) + }) +}