extproc: clear content-encoding only when body is mutated (#685)

**Commit Message**

Clearing content-type header is only necessary when the response
mutation is actually happening. Otherwise, say when the original
response is in gzip and not mutation is happening, the client gets gzip
and no content-encoding header.

**Related Issues/PRs (if applicable)**

Follow up on #683

---------

Signed-off-by: Takeshi Yoneda <[email protected]>
Co-authored-by: Erica Hughberg <[email protected]>

Author: mathetake

internal/extproc/chatcompletion_processor.go

@@ -271,16 +271,14 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 		c.metrics.RecordRequestCompletion(ctx, err == nil)
 	}()
 	var br io.Reader
-	var removeHeaders []string
+	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)
 		}
-		// 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
-		removeHeaders = append(removeHeaders, "content-encoding")
+		isGzip = true
 	default:
 		br = bytes.NewReader(body.Body)
 	}
@@ -289,10 +287,18 @@ func (c *chatCompletionProcessorUpstreamFilter) ProcessResponseBody(ctx context.
 	if err != nil {
 		return nil, fmt.Errorf("failed to transform response: %w", err)
 	}
-	if headerMutation == nil {
-		headerMutation = &extprocv3.HeaderMutation{}
+	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")
 	}
-	headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, removeHeaders...)
 
 	resp := &extprocv3.ProcessingResponse{
 		Response: &extprocv3.ProcessingResponse_ResponseBody{

tests/extproc/testupstream_test.go

@@ -164,6 +164,18 @@ func TestWithTestUpstream(t *testing.T) {
 			expStatus:       http.StatusOK,
 			expResponseBody: `{"choices":[{"message":{"content":"This is a test."}}]}`,
 		},
+		{
+			name:            "openai - /v1/chat/completions - gzip",
+			backend:         "openai",
+			responseType:    "gzip",
+			path:            "/v1/chat/completions",
+			method:          http.MethodPost,
+			requestBody:     `{"model":"something","messages":[{"role":"system","content":"You are a chatbot."}]}`,
+			expPath:         "/v1/chat/completions",
+			responseBody:    `{"choices":[{"message":{"content":"This is a test."}}]}`,
+			expStatus:       http.StatusOK,
+			expResponseBody: `{"choices":[{"message":{"content":"This is a test."}}]}`,
+		},
 		{
 			name:            "azure-openai - /v1/chat/completions",
 			backend:         "azure-openai",

tests/internal/testupstreamlib/testupstream.go

@@ -12,6 +12,7 @@ const (
 	// 	Each line in x-response-body is treated as a separate [data] payload.
 	//	* If this is "aws-event-stream", the response body is expected to be an AWS Event Stream.
 	// 	Each line in x-response-body is treated as a separate event payload.
+	//	* If this is gzip, the response body is expected to be a regular JSON response compressed with gzip.
 	//	* If this is empty, the response body is expected to be a regular JSON response.
 	ResponseTypeKey = "x-response-type"
 	// ExpectedHeadersKey is the key for the expected headers in the request.

tests/internal/testupstreamlib/testupstream/main.go

@@ -7,6 +7,7 @@ package main
 
 import (
 	"bytes"
+	"compress/gzip"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -287,8 +288,11 @@ func handler(w http.ResponseWriter, r *http.Request) {
 		logger.Println("response sent")
 		r.Context().Done()
 	default:
-		w.Header().Set("Content-Type", "application/json")
-
+		isGzip := r.Header.Get(testupstreamlib.ResponseTypeKey) == "gzip"
+		if isGzip {
+			w.Header().Set("content-encoding", "gzip")
+		}
+		w.Header().Set("content-type", "application/json")
 		var responseBody []byte
 		if expResponseBody := r.Header.Get(testupstreamlib.ResponseBodyHeaderKey); expResponseBody == "" {
 			// If the expected response body is not set, get the fake response if the path is known.
@@ -308,6 +312,13 @@ func handler(w http.ResponseWriter, r *http.Request) {
 		}
 
 		w.WriteHeader(status)
+		if isGzip {
+			var buf bytes.Buffer
+			gz := gzip.NewWriter(&buf)
+			_, _ = gz.Write(responseBody)
+			_ = gz.Close()
+			responseBody = buf.Bytes()
+		}
 		_, _ = w.Write(responseBody)
 		logger.Println("response sent:", string(responseBody))
 	}

tests/internal/testupstreamlib/testupstream/main_test.go

@@ -218,30 +218,37 @@ func Test_main(t *testing.T) {
 
 	t.Run("fake response", func(t *testing.T) {
 		t.Parallel()
-		request, err := http.NewRequest("GET",
-			"http://"+l.Addr().String()+"/v1/chat/completions", bytes.NewBuffer([]byte("expected request body")))
-		require.NoError(t, err)
+		for _, isGzip := range []bool{false, true} {
+			t.Run(fmt.Sprintf("gzip=%t", isGzip), func(t *testing.T) {
+				request, err := http.NewRequest("GET",
+					"http://"+l.Addr().String()+"/v1/chat/completions", bytes.NewBuffer([]byte("expected request body")))
+				require.NoError(t, err)
 
-		request.Header.Set(testupstreamlib.ExpectedPathHeaderKey,
-			base64.StdEncoding.EncodeToString([]byte("/v1/chat/completions")))
-		request.Header.Set(testupstreamlib.ExpectedRequestBodyHeaderKey,
-			base64.StdEncoding.EncodeToString([]byte("expected request body")))
+				request.Header.Set(testupstreamlib.ExpectedPathHeaderKey,
+					base64.StdEncoding.EncodeToString([]byte("/v1/chat/completions")))
+				request.Header.Set(testupstreamlib.ExpectedRequestBodyHeaderKey,
+					base64.StdEncoding.EncodeToString([]byte("expected request body")))
+				if isGzip {
+					request.Header.Set(testupstreamlib.ResponseTypeKey, "gzip")
+				}
 
-		response, err := http.DefaultClient.Do(request)
-		require.NoError(t, err)
-		defer func() {
-			_ = response.Body.Close()
-		}()
+				response, err := http.DefaultClient.Do(request)
+				require.NoError(t, err)
+				defer func() {
+					_ = response.Body.Close()
+				}()
 
-		require.Equal(t, http.StatusOK, response.StatusCode)
+				require.Equal(t, http.StatusOK, response.StatusCode)
 
-		responseBody, err := io.ReadAll(response.Body)
-		require.NoError(t, err)
+				responseBody, err := io.ReadAll(response.Body)
+				require.NoError(t, err)
 
-		var chat openai.ChatCompletion
-		require.NoError(t, chat.UnmarshalJSON(responseBody))
-		// Ensure that the response is one of the fake responses.
-		require.Contains(t, chatCompletionFakeResponses, chat.Choices[0].Message.Content)
+				var chat openai.ChatCompletion
+				require.NoError(t, chat.UnmarshalJSON(responseBody))
+				// Ensure that the response is one of the fake responses.
+				require.Contains(t, chatCompletionFakeResponses, chat.Choices[0].Message.Content)
+			})
+		}
 	})
 
 	t.Run("fake response for unknown path", func(t *testing.T) {