fix: reset previously set/removed values in per-backend header mutation

internal/extproc/headermutator/header_mutator.go

@@ -12,6 +12,7 @@ import (
 	extprocv3 "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 
 	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
 type HeaderMutator struct {
@@ -62,8 +63,8 @@ func (h *HeaderMutator) Mutate(headers map[string]string, onRetry bool) *extproc
 		}
 	}
 
-	// Restore original headers on retry, only if not being removed, set or not already present.
-	if onRetry && h.originalHeaders != nil {
+	if onRetry {
+		// Restore original headers on retry, only if not being removed, set or not already present.
 		for h, v := range h.originalHeaders {
 			key := strings.ToLower(h)
 			_, isRemoved := removedHeadersSet[key]
@@ -76,7 +77,25 @@ func (h *HeaderMutator) Mutate(headers map[string]string, onRetry bool) *extproc
 				})
 			}
 		}
+		// 1. Remove any headers that were added in the previous attempt (not part of original headers and not being set now)
+		// 2. Restore any original headers that were modified in the previous attempt (and not being set now).
+		for key := range headers {
+			key = strings.ToLower(key)
+			originalValue, exists := h.originalHeaders[key]
+			_, set := setHeadersSet[key]
+			if !set && !strings.HasPrefix(key, internalapi.EnvoyAIGatewayHeaderPrefix) {
+				if !exists {
+					delete(headers, key)
+					headerMutation.RemoveHeaders = append(headerMutation.RemoveHeaders, key)
+				} else {
+					// Restore original value.
+					headers[key] = originalValue
+					headerMutation.SetHeaders = append(headerMutation.SetHeaders, &corev3.HeaderValueOption{
+						Header: &corev3.HeaderValue{Key: key, RawValue: []byte(originalValue)},
+					})
+				}
+			}
+		}
 	}
-
 	return headerMutation
 }

internal/extproc/headermutator/header_mutator_test.go

@@ -41,13 +41,17 @@ func TestHeaderMutator_Mutate(t *testing.T) {
 
 	t.Run("restore original headers on retry", func(t *testing.T) {
 		originalHeaders := map[string]string{
-			"authorization": "secret",
-			"x-api-key":     "key123",
-			"other":         "value",
+			"authorization":    "secret",
+			"x-api-key":        "key123",
+			"other":            "value",
+			"only-in-original": "original",
+			"in-original-too-but-previous-attempt-set": "pikachu",
 		}
 		headers := map[string]string{
 			"other":         "value",
 			"authorization": "secret",
+			"in-original-too-but-previous-attempt-set": "charmander",
+			"only-set-previously":                      "bulbasaur",
 		}
 		mutations := &filterapi.HTTPHeaderMutation{
 			Remove: []string{"authorization"},
@@ -57,9 +61,11 @@ func TestHeaderMutator_Mutate(t *testing.T) {
 		mutation := mutator.Mutate(headers, true)
 
 		require.NotNil(t, mutation)
-		require.ElementsMatch(t, []string{"authorization"}, mutation.RemoveHeaders)
+		require.ElementsMatch(t, []string{"authorization", "only-set-previously"}, mutation.RemoveHeaders)
 		require.Equal(t, "key123", headers["x-api-key"])
 		require.Equal(t, "value", headers["other"])
 		require.Equal(t, "secret", headers["authorization"])
+		require.Equal(t, "original", headers["only-in-original"])
+		require.Equal(t, "pikachu", headers["in-original-too-but-previous-attempt-set"])
 	})
 }

internal/extproc/server.go

@@ -132,11 +132,11 @@ func (s *Server) processorForPath(requestHeaders map[string]string, isUpstreamFi
 
 // originalPathHeader is the header used to pass the original path to the processor.
 // This is used in the upstream filter level to determine the original path of the request on retry.
-const originalPathHeader = "x-ai-eg-original-path"
+const originalPathHeader = internalapi.EnvoyAIGatewayHeaderPrefix + "original-path"
 
 // internalReqIDHeader is the header used to pass the unique internal request ID to the upstream filter.
 // This ensures that the upstream filter uses the same unique ID as the router filter to avoid race conditions.
-const internalReqIDHeader = "x-ai-eg-internal-req-id"
+const internalReqIDHeader = internalapi.EnvoyAIGatewayHeaderPrefix + "internal-req-id"
 
 // Process implements [extprocv3.ExternalProcessorServer].
 func (s *Server) Process(stream extprocv3.ExternalProcessor_ProcessServer) error {

internal/internalapi/internalapi.go

@@ -15,14 +15,16 @@ import (
 )
 
 const (
+	// EnvoyAIGatewayHeaderPrefix is the prefix for special headers used by AI Gateway, either for internal or external use.
+	EnvoyAIGatewayHeaderPrefix = "x-ai-eg-"
 	// InternalEndpointMetadataNamespace is the namespace used for the dynamic metadata for internal use.
 	InternalEndpointMetadataNamespace = "aigateway.envoy.io"
 	// InternalMetadataBackendNameKey is the key used to store the backend name
 	InternalMetadataBackendNameKey = "per_route_rule_backend_name"
 	// MCPBackendHeader is the special header key used to specify the target backend name.
-	MCPBackendHeader = "x-ai-eg-mcp-backend"
+	MCPBackendHeader = EnvoyAIGatewayHeaderPrefix + "eg-mcp-backend"
 	// MCPRouteHeader is the special header key used to identify the mcp route.
-	MCPRouteHeader = "x-ai-eg-mcp-route"
+	MCPRouteHeader = EnvoyAIGatewayHeaderPrefix + "mcp-route"
 	// MCPBackendListenerPort is the port for the MCP backend listener.
 	MCPBackendListenerPort = 10088
 	// MCPProxyPort is the port where the MCP proxy listens.

tests/extproc/envoy.yaml

@@ -461,7 +461,7 @@ static_resources:
                 metadata:
                   filter_metadata:
                     aigateway.envoy.io:
-                      per_route_rule_backend_name: "testupstream-openai-always-200"
+                      per_route_rule_backend_name: "testupstream-openai"
             priority: 1 # Secondary.
     - name: testupstream-modelname-override
       connect_timeout: 0.25s

tests/extproc/testupstream_test.go

@@ -61,10 +61,6 @@ func TestWithTestUpstream(t *testing.T) {
 			testUpstreamAzureBackend,
 			testUpstreamGCPVertexAIBackend,
 			testUpstreamGCPAnthropicAIBackend,
-			// TODO: this shouldn't be needed. The previous per-backend headers shouldn't affect the subsequent retries.
-			{Name: "testupstream-openai-always-200", Schema: openAISchema, HeaderMutation: &filterapi.HTTPHeaderMutation{
-				Set: []filterapi.HTTPHeader{{Name: testupstreamlib.ResponseStatusKey, Value: "200"}},
-			}},
 			{
 				Name: "testupstream-openai-5xx", Schema: openAISchema, HeaderMutation: &filterapi.HTTPHeaderMutation{
 					Set: []filterapi.HTTPHeader{{Name: testupstreamlib.ResponseStatusKey, Value: "500"}},

tests/internal/testenvironment/test_environment.go

@@ -262,8 +262,8 @@ func requireEnvoy(t testing.TB,
 		"--concurrency", strconv.Itoa(max(runtime.NumCPU(), 2)),
 		// This allows multiple Envoy instances to run in parallel.
 		"--base-id", strconv.Itoa(time.Now().Nanosecond()),
-		// Add debug logging for extproc.
-		"--component-log-level", "ext_proc:trace,http:debug,connection:debug",
+		// Add debug logging for http.
+		"--component-log-level", "http:debug",
 	)
 
 	// wait for the ready message or exit.