From ce71608d5382131d1c9b42b28a8884271c5806a8 Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Thu, 20 Aug 2026 14:44:46 +0530 Subject: [PATCH] fix: support compressed Anthropic OAuth passthrough respones --- core/providers/anthropic/anthropic.go | 4 +- .../anthropic/passthroughheaders_test.go | 69 +++++++- core/providers/utils/utils.go | 164 +++++++++++++++--- core/providers/utils/utils_test.go | 61 ++++++- pulse.yaml | 1 + ui/app/workspace/logs/page.tsx | 18 +- 6 files changed, 274 insertions(+), 43 deletions(-) diff --git a/core/providers/anthropic/anthropic.go b/core/providers/anthropic/anthropic.go index 598f928d972..7355e142601 100644 --- a/core/providers/anthropic/anthropic.go +++ b/core/providers/anthropic/anthropic.go @@ -778,7 +778,7 @@ func HandleAnthropicChatCompletionStreaming( providerUtils.SetExtraHeaders(ctx, req, extraHeaders, []string{AnthropicBetaHeader}) // OAuth passthrough: forward the caller's raw headers, whose token is the upstream // credential. Skips anthropic-beta — MergeBetaHeaders below owns the final value. - providerUtils.SetPassthroughHeaders(ctx, req, providerName, []string{AnthropicBetaHeader}) + providerUtils.SetPassthroughHeadersForStreaming(ctx, req, providerName, []string{AnthropicBetaHeader}) if betaHeaders := FilterBetaHeadersForProvider(MergeBetaHeaders(ctx, extraHeaders), providerName, betaHeaderOverrides); len(betaHeaders) > 0 { req.Header.Set(AnthropicBetaHeader, strings.Join(betaHeaders, ",")) @@ -1420,7 +1420,7 @@ func HandleAnthropicResponsesStream( providerUtils.SetExtraHeaders(ctx, req, extraHeaders, []string{AnthropicBetaHeader}) // OAuth passthrough: forward the caller's raw headers, whose token is the upstream // credential. Skips anthropic-beta — MergeBetaHeaders below owns the final value. - providerUtils.SetPassthroughHeaders(ctx, req, providerName, []string{AnthropicBetaHeader}) + providerUtils.SetPassthroughHeadersForStreaming(ctx, req, providerName, []string{AnthropicBetaHeader}) if betaHeaders := FilterBetaHeadersForProvider(MergeBetaHeaders(ctx, extraHeaders), providerName, betaHeaderOverrides); len(betaHeaders) > 0 { req.Header.Set(AnthropicBetaHeader, strings.Join(betaHeaders, ",")) diff --git a/core/providers/anthropic/passthroughheaders_test.go b/core/providers/anthropic/passthroughheaders_test.go index 1258f336d6f..762827cf983 100644 --- a/core/providers/anthropic/passthroughheaders_test.go +++ b/core/providers/anthropic/passthroughheaders_test.go @@ -19,6 +19,7 @@ func clientHeaders() map[string][]string { "anthropic-beta": {"interleaved-thinking-2025-05-14"}, "anthropic-version": {"2023-06-01"}, "content-type": {"text/plain"}, + "accept-encoding": {"gzip, deflate, br"}, "x-app": {"cli"}, "x-stainless-lang": {"js"}, "x-bf-vk": {"sk-bf-must-never-leave-the-gateway"}, @@ -41,6 +42,7 @@ func TestSetPassthroughHeaders_ForwardsCallerCredentialButNotInternals(t *testin forwarded := map[string]string{ "authorization": "Bearer sk-ant-oat01-caller-token", "anthropic-version": "2023-06-01", + "accept-encoding": "gzip, deflate, br", "x-app": "cli", "x-stainless-lang": "js", } @@ -59,10 +61,69 @@ func TestSetPassthroughHeaders_ForwardsCallerCredentialButNotInternals(t *testin if got := string(req.Header.Peek(AnthropicBetaHeader)); got != "" { t.Errorf("anthropic-beta must be left to MergeBetaHeaders, got %q", got) } - // Bifrost owns the body, so it owns Content-Type — a caller value must never override it, - // regardless of where callers invoke this relative to their own SetContentType. - if got := string(req.Header.Peek("Content-Type")); got != "" { - t.Errorf("caller Content-Type must not be forwarded, got %q", got) + // Bifrost owns the request body, so the caller cannot override its content type. + if got := string(req.Header.Peek("content-type")); got != "" { + t.Errorf("caller content-type must not be forwarded, got %q", got) + } +} + +func TestSetPassthroughHeaders_FiltersAcceptEncodingToSupportedCodecs(t *testing.T) { + tests := []struct { + name string + values []string + streaming bool + wantForward string + }{ + { + name: "supported list is preserved", + values: []string{"gzip, deflate, br, zstd"}, + wantForward: "gzip, deflate, br, zstd", + }, + { + name: "unsupported tokens and wildcard are removed", + values: []string{"gzip, snappy;q=0.8", "br;q=0.5, *;q=0"}, + wantForward: "gzip, br;q=0.5", + }, + { + name: "unsupported-only header pins identity", + values: []string{"snappy"}, + wantForward: "identity", + }, + { + name: "streaming keeps only incrementally decoded codecs", + values: []string{"gzip, deflate, br, zstd"}, + streaming: true, + wantForward: "gzip", + }, + { + // Dropping the header here would let the upstream answer in br, which the + // streaming decoder cannot handle. + name: "streaming pins identity when no codec survives", + values: []string{"br, zstd"}, + streaming: true, + wantForward: "identity", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + ctx.SetValue(schemas.BifrostContextKeyPassthroughHeaders, map[string][]string{ + "accept-encoding": tt.values, + }) + + req := fasthttp.AcquireRequest() + defer fasthttp.ReleaseRequest(req) + if tt.streaming { + providerUtils.SetPassthroughHeadersForStreaming(ctx, req, schemas.Anthropic, nil) + } else { + providerUtils.SetPassthroughHeaders(ctx, req, schemas.Anthropic, nil) + } + + if got := string(req.Header.Peek("accept-encoding")); got != tt.wantForward { + t.Errorf("Accept-Encoding = %q, want %q", got, tt.wantForward) + } + }) } } diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index 848212e1435..b51e05301b8 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -999,12 +999,68 @@ func SetExtraHeaders(ctx context.Context, req *fasthttp.Request, extraHeaders ma // gateway credentials and must never be forwarded to a provider. const internalHeaderPrefix = "x-bf-" +// supportedBufferedContentEncodings are the response codings handled by +// CheckAndDecodeBody. Accept-Encoding passthrough is intersected with this set +// so Bifrost never advertises an upstream response format it cannot parse. +var supportedBufferedContentEncodings = map[string]struct{}{ + "identity": {}, + "gzip": {}, + "x-gzip": {}, + "deflate": {}, + "br": {}, + "zstd": {}, +} + +// supportedStreamingContentEncodings are the codings handled incrementally by +// DecompressStreamBody. Keep this separate from the buffered set: advertising +// Brotli or zstd on an SSE request would be unsafe until the stream reader can +// decode those formats without buffering the whole response. +var supportedStreamingContentEncodings = map[string]struct{}{ + "identity": {}, + "gzip": {}, + "x-gzip": {}, +} + +func filterSupportedAcceptEncodings(values []string, supportedEncodings map[string]struct{}) []string { + filtered := make([]string, 0, len(values)) + for _, value := range values { + for item := range strings.SplitSeq(value, ",") { + item = strings.TrimSpace(item) + if item == "" { + continue + } + + coding := item + if separator := strings.IndexByte(coding, ';'); separator >= 0 { + coding = coding[:separator] + } + coding = strings.ToLower(strings.TrimSpace(coding)) + if _, ok := supportedEncodings[coding]; ok { + // Preserve parameters such as q-values while normalizing the list into + // one header value below. + filtered = append(filtered, item) + } + } + } + return filtered +} + // SetPassthroughHeaders applies the caller's raw request headers captured for Anthropic OAuth // passthrough, where the caller's token is the upstream credential. ONLY the Anthropic provider // may call this: every other provider authenticates with its own configured credentials, so // forwarding these would leak x-bf-* upstream and, on Bedrock, break SigV4 when a hop rewrites // x-forwarded-for. Hop-by-hop headers are dropped by filterHeaders, x-bf-* never leaves the gateway. func SetPassthroughHeaders(ctx context.Context, req *fasthttp.Request, provider schemas.ModelProvider, skipHeaders []string) { + setPassthroughHeaders(ctx, req, provider, skipHeaders, supportedBufferedContentEncodings) +} + +// SetPassthroughHeadersForStreaming applies OAuth passthrough headers while +// advertising only response codings that Bifrost can decode incrementally. +func SetPassthroughHeadersForStreaming(ctx context.Context, req *fasthttp.Request, provider schemas.ModelProvider, skipHeaders []string) { + setPassthroughHeaders(ctx, req, provider, skipHeaders, supportedStreamingContentEncodings) +} + +func setPassthroughHeaders(ctx context.Context, req *fasthttp.Request, provider schemas.ModelProvider, skipHeaders []string, supportedEncodings map[string]struct{}) { // Gate on the provider here rather than at the call sites: the Anthropic request handlers // are shared with azure, vertex, bedrockmantle, vllm, sgl, deepseek and fireworks, so a // call-site check would silently forward the caller's credential to all of them. @@ -1020,13 +1076,25 @@ func SetPassthroughHeaders(ctx context.Context, req *fasthttp.Request, provider if strings.HasPrefix(lower, internalHeaderPrefix) { continue } - // Bifrost owns the body it sends, so it owns Content-Type. Dropping it here keeps a - // caller value from overriding the JSON content type no matter where callers invoke - // this relative to their own SetContentType. + if skipHeaders != nil && slices.Contains(skipHeaders, lower) { + continue + } + // Bifrost owns the body it sends, so a caller cannot override its content type. if lower == "content-type" { continue } - if skipHeaders != nil && slices.Contains(skipHeaders, lower) { + if lower == "accept-encoding" { + // Filtering must never widen what the upstream may send: an omitted + // Accept-Encoding means any coding is acceptable (RFC 9110 12.5.3), so a + // caller list that filters down to nothing pins identity instead of + // dropping the header. Without this, a streaming caller asking for br + // alone would let the upstream answer in br, which DecompressStreamBody + // cannot decode, and the SSE parser would see raw compressed bytes. + if supported := filterSupportedAcceptEncodings(values, supportedEncodings); len(supported) > 0 { + req.Header.Set(k, strings.Join(supported, ", ")) + } else { + req.Header.Set(k, "identity") + } continue } for i, v := range values { @@ -1969,33 +2037,79 @@ func CheckOperationAllowed(defaultProvider schemas.ModelProvider, config *schema } // CheckAndDecodeBody checks the content encoding and decodes the body accordingly. -// It returns a copy of the body to avoid race conditions when the response is released -// back to fasthttp's buffer pool. Uses pooled gzip readers to reduce GC pressure. +// It returns an owned body to avoid races when the response is released back to +// fasthttp's buffer pool. Content codings are decoded in reverse application order, +// as required by RFC 9110, using the shared pooled readers. func CheckAndDecodeBody(resp *fasthttp.Response) ([]byte, error) { + body := resp.Body() + if len(body) == 0 { + return nil, nil + } + + result := append([]byte(nil), body...) contentEncoding := strings.ToLower(strings.TrimSpace(string(resp.Header.Peek("Content-Encoding")))) - if strings.Contains(contentEncoding, "gzip") { - body := resp.Body() - if len(body) == 0 { - return nil, nil - } + if contentEncoding == "" || contentEncoding == "identity" { + return result, nil + } - reader := bytes.NewReader(body) - gz, err := AcquireGzipReader(reader) - if err != nil { - return nil, err - } - defer ReleaseGzipReader(gz) + encodings := strings.Split(contentEncoding, ",") + for i := len(encodings) - 1; i >= 0; i-- { + encoding := strings.TrimSpace(encodings[i]) + reader := bytes.NewReader(result) - decompressed, err := io.ReadAll(gz) - if err != nil { - return nil, err + // Release on the Acquire error paths too: zstd.NewReader's contract is to + // return the decoder alongside its error, so a bare return could drop one + // from the pool. The gzip and deflate constructors return nil there, where + // Release is a no-op. + switch encoding { + case "", "identity": + continue + case "gzip", "x-gzip": + gz, err := AcquireGzipReader(reader) + if err != nil { + ReleaseGzipReader(gz) + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + result, err = io.ReadAll(gz) + ReleaseGzipReader(gz) + if err != nil { + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + case "deflate": + fr, err := AcquireFlateReader(reader) + if err != nil { + ReleaseFlateReader(fr) + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + result, err = io.ReadAll(fr) + ReleaseFlateReader(fr) + if err != nil { + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + case "br": + br := AcquireBrotliReader(reader) + var err error + result, err = io.ReadAll(br) + ReleaseBrotliReader(br) + if err != nil { + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + case "zstd": + dec, err := AcquireZstdDecoder(reader) + if err != nil { + ReleaseZstdDecoder(dec) + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + result, err = io.ReadAll(dec) + ReleaseZstdDecoder(dec) + if err != nil { + return nil, fmt.Errorf("decode %s response body: %w", encoding, err) + } + default: + return nil, fmt.Errorf("unsupported Content-Encoding %q", encoding) } - return decompressed, nil } - // Copy the body to avoid race conditions when response is released back to pool - body := resp.Body() - result := make([]byte, len(body)) - copy(result, body) + return result, nil } diff --git a/core/providers/utils/utils_test.go b/core/providers/utils/utils_test.go index f7389480bef..27e0f832448 100644 --- a/core/providers/utils/utils_test.go +++ b/core/providers/utils/utils_test.go @@ -658,9 +658,10 @@ func TestMarshalSorted_Deterministic(t *testing.T) { } } -// TestCheckAndDecodeBody_PooledGzip verifies that CheckAndDecodeBody correctly -// decompresses gzip-encoded responses using pooled gzip readers. -func TestCheckAndDecodeBody_PooledGzip(t *testing.T) { +// TestCheckAndDecodeBody_ContentEncodings verifies that unary provider responses +// are decoded with the shared pooled readers before JSON parsing. +func TestCheckAndDecodeBody_ContentEncodings(t *testing.T) { + chainedPayload := gzipCompress([]byte(`{"message":"chained"}`)) tests := []struct { name string body []byte @@ -689,6 +690,36 @@ func TestCheckAndDecodeBody_PooledGzip(t *testing.T) { wantBody: `trimmed`, wantErr: false, }, + { + name: "brotli encoded body", + body: compressBrotli([]byte(`{"input_tokens":16998}`)), + contentEncoding: "br", + wantBody: `{"input_tokens":16998}`, + }, + { + name: "deflate encoded body", + body: compressFlate([]byte(`{"message":"deflate"}`)), + contentEncoding: "deflate", + wantBody: `{"message":"deflate"}`, + }, + { + name: "zstd encoded body", + body: compressZstd([]byte(`{"message":"zstd"}`)), + contentEncoding: "zstd", + wantBody: `{"message":"zstd"}`, + }, + { + name: "identity encoded body", + body: []byte(`{"message":"identity"}`), + contentEncoding: "identity", + wantBody: `{"message":"identity"}`, + }, + { + name: "chained gzip then brotli body", + body: compressBrotli(chainedPayload), + contentEncoding: "gzip, br", + wantBody: `{"message":"chained"}`, + }, { name: "no encoding - plain body", body: []byte(`plain text`), @@ -709,6 +740,30 @@ func TestCheckAndDecodeBody_PooledGzip(t *testing.T) { contentEncoding: "gzip", wantErr: true, }, + { + name: "invalid brotli data", + body: []byte{0xFF, 0xFE, 0xFD}, + contentEncoding: "br", + wantErr: true, + }, + { + name: "invalid deflate data", + body: []byte{0xFF, 0xFE, 0xFD}, + contentEncoding: "deflate", + wantErr: true, + }, + { + name: "invalid zstd data", + body: []byte{0xFF, 0xFE, 0xFD}, + contentEncoding: "zstd", + wantErr: true, + }, + { + name: "unsupported encoding", + body: []byte(`encoded somehow`), + contentEncoding: "snappy", + wantErr: true, + }, } for _, tt := range tests { diff --git a/pulse.yaml b/pulse.yaml index db01aadec9b..164aae603a6 100644 --- a/pulse.yaml +++ b/pulse.yaml @@ -10,6 +10,7 @@ services: kill_timeout: 10s proxy: addr: ":${PORT:-8080}" + write_timeout: 900s healthcheck: path: /health interval: 2s diff --git a/ui/app/workspace/logs/page.tsx b/ui/app/workspace/logs/page.tsx index 4ef267a16ee..b2158b8effe 100644 --- a/ui/app/workspace/logs/page.tsx +++ b/ui/app/workspace/logs/page.tsx @@ -135,20 +135,20 @@ export default function LogsPage() { cache_hit_types: urlState.cache_hit_types, metadata_filters: urlState.metadata_filters ? (() => { - try { - return JSON.parse(urlState.metadata_filters); - } catch { - return undefined; - } - })() + try { + return JSON.parse(urlState.metadata_filters); + } catch { + return undefined; + } + })() : undefined, // Use a period if present ...(urlState.period ? { period: urlState.period } : { - start_time: dateUtils.toISOString(urlState.start_time), - end_time: dateUtils.toISOString(urlState.end_time), - }), + start_time: dateUtils.toISOString(urlState.start_time), + end_time: dateUtils.toISOString(urlState.end_time), + }), }), // Only re-derive filters when filter-related URL params change (not pagination) [