fix: streamline gzip handling in response processing (envoyproxy#1189)

**Description**
Bug: Token Usage not extracted for gzipped http-responses from upstream
when endpoint is `v1/messages`

Root Cause: The `messages_processor.go` doensn't decompress response
body (if gzipped) before extracting token usage. see [1].

Corresponding code in `chatcompletion_processor.go`(See [2]) which
handles gzip decompressions.

This PR introduces utility functions for handling gzipped http-responses
and adds them to `chatcompletion_processor.go`,
`embeddings_processor.go` and `messages_processor.go`

[1]:
https://github.com/envoyproxy/ai-gateway/blob/e767da5559fadf89a01d2168eea262b48a1ab6a4/internal/extproc/messages_processor.go#L254

[2]:
https://github.com/envoyproxy/ai-gateway/blob/e767da5559fadf89a01d2168eea262b48a1ab6a4/internal/extproc/chatcompletion_processor.go#L330-L340

---------

Signed-off-by: Sukumar Gaonkar <[email protected]>

Author: sukumargaonkar

internal/extproc/chatcompletion_processor.go

@@ -6,12 +6,9 @@
 package extproc
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log/slog"
 	"strconv"
 
@@ -326,24 +323,18 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 			c.metrics.RecordRequestCompletion(ctx, true, c.requestHeaders)
 		}
 	}()
-	var br io.Reader
-	var isGzip bool
-	switch c.responseEncoding {
-	case "gzip":
-		br, err = gzip.NewReader(bytes.NewReader(body.Body))
-		if err != nil {
-			return nil, fmt.Errorf("failed to decode gzip: %w", err)
-		}
-		isGzip = true
-	default:
-		br = bytes.NewReader(body.Body)
+
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, c.responseEncoding)
+	if err != nil {
+		return nil, err
 	}
 
 	// Assume all responses have a valid status code header.
 	if code, _ := strconv.Atoi(c.responseHeaders[":status"]); !isGoodStatusCode(code) {
 		var headerMutation *extprocv3.HeaderMutation
 		var bodyMutation *extprocv3.BodyMutation
-		headerMutation, bodyMutation, err = c.translator.ResponseError(c.responseHeaders, br)
+		headerMutation, bodyMutation, err = c.translator.ResponseError(c.responseHeaders, decodingResult.reader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to transform response error: %w", err)
 		}
@@ -368,22 +359,13 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 		}, nil
 	}
 
-	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream, c.span)
+	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, decodingResult.reader, body.EndOfStream, c.span)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
-	if bodyMutation != nil && isGzip {
-		if headerMutation == nil {
-			headerMutation = &extprocv3.HeaderMutation{}
-		}
-		// TODO: this is a hotfix, we should update this to recompress since its in the header
-		// If the response was gzipped, ensure we remove the content-encoding header.
-		//
-		// This is only needed when the transformation is actually modifying the body. When the backend
-		// is in OpenAI format (and it's the first try before any retry), the response body is not modified,
-		// so we don't need to remove the header in that case.
-		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
-	}
+
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
 
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{

internal/extproc/embeddings_processor.go

@@ -6,12 +6,9 @@
 package extproc
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
-	"io"
 	"log/slog"
 	"strconv"
 
@@ -250,24 +247,18 @@ func (e *embeddingsProcessorUpstreamFilter) ProcessResponseBody(ctx context.Cont
 			e.metrics.RecordRequestCompletion(ctx, true, e.requestHeaders)
 		}
 	}()
-	var br io.Reader
-	var isGzip bool
-	switch e.responseEncoding {
-	case "gzip":
-		br, err = gzip.NewReader(bytes.NewReader(body.Body))
-		if err != nil {
-			return nil, fmt.Errorf("failed to decode gzip: %w", err)
-		}
-		isGzip = true
-	default:
-		br = bytes.NewReader(body.Body)
+
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, e.responseEncoding)
+	if err != nil {
+		return nil, err
 	}
 
 	// Assume all responses have a valid status code header.
 	if code, _ := strconv.Atoi(e.responseHeaders[":status"]); !isGoodStatusCode(code) {
 		var headerMutation *extprocv3.HeaderMutation
 		var bodyMutation *extprocv3.BodyMutation
-		headerMutation, bodyMutation, err = e.translator.ResponseError(e.responseHeaders, br)
+		headerMutation, bodyMutation, err = e.translator.ResponseError(e.responseHeaders, decodingResult.reader)
 		if err != nil {
 			return nil, fmt.Errorf("failed to transform response error: %w", err)
 		}
@@ -285,22 +276,13 @@ func (e *embeddingsProcessorUpstreamFilter) ProcessResponseBody(ctx context.Cont
 		}, nil
 	}
 
-	headerMutation, bodyMutation, tokenUsage, err := e.translator.ResponseBody(e.responseHeaders, br, body.EndOfStream)
+	headerMutation, bodyMutation, tokenUsage, err := e.translator.ResponseBody(e.responseHeaders, decodingResult.reader, body.EndOfStream)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
-	if bodyMutation != nil && isGzip {
-		if headerMutation == nil {
-			headerMutation = &extprocv3.HeaderMutation{}
-		}
-		// TODO: this is a hotfix, we should update this to recompress since its in the header
-		// If the response was gzipped, ensure we remove the content-encoding header.
-		//
-		// This is only needed when the transformation is actually modifying the body. When the backend
-		// is in OpenAI format (and it's the first try before any retry), the response body is not modified,
-		// so we don't need to remove the header in that case.
-		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
-	}
+
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
 
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{

internal/extproc/messages_processor.go

@@ -6,7 +6,6 @@
 package extproc
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -134,6 +133,7 @@ type messagesProcessorUpstreamFilter struct {
 	config                 *processorConfig
 	requestHeaders         map[string]string
 	responseHeaders        map[string]string
+	responseEncoding       string
 	modelNameOverride      string
 	backendName            string
 	handler                backendauth.Handler
@@ -222,6 +222,9 @@ func (c *messagesProcessorUpstreamFilter) ProcessResponseHeaders(ctx context.Con
 	}()
 
 	c.responseHeaders = headersToMap(headers)
+	if enc := c.responseHeaders["content-encoding"]; enc != "" {
+		c.responseEncoding = enc
+	}
 	headerMutation, err := c.translator.ResponseHeaders(c.responseHeaders)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response headers: %w", err)
@@ -250,14 +253,21 @@ func (c *messagesProcessorUpstreamFilter) ProcessResponseBody(ctx context.Contex
 		}
 	}()
 
-	// Simple passthrough: just pass the body as-is without any complex handling.
-	br := bytes.NewReader(body.Body)
+	// Decompress the body if needed using common utility.
+	decodingResult, err := decodeContentIfNeeded(body.Body, c.responseEncoding)
+	if err != nil {
+		return nil, err
+	}
 
-	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream)
+	// headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, br, body.EndOfStream).
+	headerMutation, bodyMutation, tokenUsage, err := c.translator.ResponseBody(c.responseHeaders, decodingResult.reader, body.EndOfStream)
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
 
+	// Remove content-encoding header if original body encoded but was mutated in the processor.
+	headerMutation = removeContentEncodingIfNeeded(headerMutation, bodyMutation, decodingResult.isEncoded)
+
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{
 			ResponseBody: &extprocv3.BodyResponse{

internal/extproc/util.go

@@ -5,6 +5,60 @@
 
 package extproc
 
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
+
+	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
+)
+
+// contentDecodingResult contains the result of content decoding operation.
+type contentDecodingResult struct {
+	reader    io.Reader
+	isEncoded bool
+}
+
+// decodeContentIfNeeded decompresses the response body based on the content-encoding header.
+// Currently supports gzip encoding, but can be extended to support other encodings in the future.
+// Returns a reader for the (potentially decompressed) body and metadata about the encoding.
+func decodeContentIfNeeded(body []byte, contentEncoding string) (contentDecodingResult, error) {
+	switch contentEncoding {
+	case "gzip":
+		reader, err := gzip.NewReader(bytes.NewReader(body))
+		if err != nil {
+			return contentDecodingResult{}, fmt.Errorf("failed to decode gzip: %w", err)
+		}
+		return contentDecodingResult{
+			reader:    reader,
+			isEncoded: true,
+		}, nil
+	default:
+		return contentDecodingResult{
+			reader:    bytes.NewReader(body),
+			isEncoded: false,
+		}, nil
+	}
+}
+
+// removeContentEncodingIfNeeded removes the content-encoding header if the body was modified and was encoded.
+// This is needed when the transformation modifies the body content but the response was originally compressed.
+func removeContentEncodingIfNeeded(headerMutation *extprocv3.HeaderMutation, bodyMutation *extprocv3.BodyMutation, isEncoded bool) *extprocv3.HeaderMutation {
+	if bodyMutation != nil && isEncoded {
+		if headerMutation == nil {
+			headerMutation = &extprocv3.HeaderMutation{}
+		}
+		// TODO: this is a hotfix, we should update this to recompress since its in the header
+		// If the upstream response was compressed and we decompressed it,
+		// ensure we remove the content-encoding header.
+		//
+		// This is only needed when the transformation is actually modifying the body.
+		headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, "content-encoding")
+	}
+	return headerMutation
+}
+
 // isGoodStatusCode checks if the HTTP status code of the upstream response is successful.
 // The 2xx - Successful: The request is received by upstream and processed successfully.
 // https://developer.mozilla.org/en-US/docs/Web/HTTP/Status#successful_responses

internal/extproc/util_test.go

@@ -6,8 +6,12 @@
 package extproc
 
 import (
+	"bytes"
+	"compress/gzip"
+	"io"
 	"testing"
 
+	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"github.com/stretchr/testify/require"
 )
 
@@ -19,3 +23,123 @@ func TestIsGoodStatusCode(t *testing.T) {
 		require.False(t, isGoodStatusCode(s))
 	}
 }
+
+func TestDecodeContentIfNeeded(t *testing.T) {
+	tests := []struct {
+		name         string
+		body         []byte
+		encoding     string
+		wantEncoded  bool
+		wantEncoding string
+		wantErr      bool
+	}{
+		{
+			name:         "plain body",
+			body:         []byte("hello world"),
+			encoding:     "",
+			wantEncoded:  false,
+			wantEncoding: "",
+			wantErr:      false,
+		},
+		{
+			name:         "unsupported encoding",
+			body:         []byte("hello world"),
+			encoding:     "deflate",
+			wantEncoded:  false,
+			wantEncoding: "",
+			wantErr:      false,
+		},
+		{
+			name: "valid gzip",
+			body: func() []byte {
+				var b bytes.Buffer
+				w := gzip.NewWriter(&b)
+				_, err := w.Write([]byte("abc"))
+				if err != nil {
+					panic(err)
+				}
+				w.Close()
+				return b.Bytes()
+			}(),
+			encoding:     "gzip",
+			wantEncoded:  true,
+			wantEncoding: "gzip",
+			wantErr:      false,
+		},
+		{
+			name:         "invalid gzip",
+			body:         []byte("not a gzip"),
+			encoding:     "gzip",
+			wantEncoded:  false,
+			wantEncoding: "",
+			wantErr:      true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			res, err := decodeContentIfNeeded(tt.body, tt.encoding)
+			if tt.wantErr {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, tt.wantEncoded, res.isEncoded)
+			if !tt.wantEncoded {
+				out, _ := io.ReadAll(res.reader)
+				require.Equal(t, tt.body, out)
+			} else if tt.encoding == "gzip" && !tt.wantErr {
+				out, _ := io.ReadAll(res.reader)
+				require.Equal(t, []byte("abc"), out)
+			}
+		})
+	}
+}
+
+func TestRemoveContentEncodingIfNeeded(t *testing.T) {
+	tests := []struct {
+		name        string
+		hm          *extprocv3.HeaderMutation
+		bm          *extprocv3.BodyMutation
+		isEncoded   bool
+		wantRemoved bool
+	}{
+		{
+			name:        "no body mutation, not encoded",
+			hm:          nil,
+			bm:          nil,
+			isEncoded:   false,
+			wantRemoved: false,
+		},
+		{
+			name:        "body mutation, not encoded",
+			hm:          nil,
+			bm:          &extprocv3.BodyMutation{},
+			isEncoded:   false,
+			wantRemoved: false,
+		},
+		{
+			name:        "body mutation, encoded",
+			hm:          nil,
+			bm:          &extprocv3.BodyMutation{},
+			isEncoded:   true,
+			wantRemoved: true,
+		},
+		{
+			name:        "existing header mutation, body mutation, encoded",
+			hm:          &extprocv3.HeaderMutation{RemoveHeaders: []string{"foo"}},
+			bm:          &extprocv3.BodyMutation{},
+			isEncoded:   true,
+			wantRemoved: true,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			res := removeContentEncodingIfNeeded(tt.hm, tt.bm, tt.isEncoded)
+			if tt.wantRemoved {
+				require.Contains(t, res.RemoveHeaders, "content-encoding")
+			} else if res != nil {
+				require.NotContains(t, res.RemoveHeaders, "content-encoding")
+			}
+		})
+	}
+}