From cd9d24c99ae41b3fd0cd00dbb6a69a405fc62bee Mon Sep 17 00:00:00 2001 From: Saylee Date: Thu, 5 Nov 2026 18:21:04 +0530 Subject: [PATCH] fix(router): drop a failed attempt's response headers before retrying proxyRequest stages the backend's status and headers before forwarding the body. When a backend sends a complete header block and closes before sending the declared body, ForwardBody returns an error without writing anything, so the retry loop continues on another pod. The failed attempt's headers are still staged, and copyResponseHeaders appends the healthy attempt's values. The client then receives two Content-Length headers plus stale cookies and headers from the failed backend. Snapshot the headers and status that earlier handlers staged, and restore them when forwarding fails before anything is committed. Headers added by router middleware survive, and the healthy retry carries one Content-Length and only its own upstream headers. TestProxy_RetryDropsFailedAttemptHeaders drives the production path with a real backend that declares 4096 bytes and closes, followed by a healthy retry. Signed-off-by: Saylee --- pkg/kthena-router/router/router.go | 34 ++++++- pkg/kthena-router/router/router_test.go | 113 ++++++++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) diff --git a/pkg/kthena-router/router/router.go b/pkg/kthena-router/router/router.go index 698b9e0708..47d1736448 100644 --- a/pkg/kthena-router/router/router.go +++ b/pkg/kthena-router/router/router.go @@ -1282,7 +1282,18 @@ func proxyRequest( path = "/v1/responses" } parser := providers.DefaultAdapter().ResponseParser(c, path) - return forwardResponseWithUsageParser(c, resp, stream, parser, onUsage) + + // The caller may retry another backend while nothing has been written + // downstream. Snapshot what earlier handlers staged so headers copied from a + // backend that fails before sending its body cannot leak into the retry. + stagedHeaders := cloneResponseHeaders(c.Writer.Header()) + stagedStatus := c.Writer.Status() + forwardErr := forwardResponseWithUsageParser(c, resp, stream, parser, onUsage) + if forwardErr != nil && !c.Writer.Written() { + restoreResponseHeaders(c.Writer.Header(), stagedHeaders) + c.Status(stagedStatus) + } + return forwardErr } // originalRequestPath returns the client-facing request path used for wire @@ -1390,6 +1401,27 @@ func copyResponseHeaders(c *gin.Context, headers http.Header, stream bool) { } } +// cloneResponseHeaders returns a copy of headers so the values stay unchanged +// after c.Writer.Header() is modified. +func cloneResponseHeaders(headers http.Header) http.Header { + clone := make(http.Header, len(headers)) + for key, values := range headers { + clone[key] = append([]string(nil), values...) + } + return clone +} + +// restoreResponseHeaders replaces the staged downstream headers with a copy of +// headers, dropping the entries a failed upstream attempt added. +func restoreResponseHeaders(dst http.Header, headers http.Header) { + for key := range dst { + delete(dst, key) + } + for key, values := range headers { + dst[key] = append([]string(nil), values...) + } +} + func shouldSkipResponseHeader(header string, stream bool, dynamicHopByHopHeaders map[string]struct{}) bool { if stream && strings.EqualFold(header, "Content-Length") { return true diff --git a/pkg/kthena-router/router/router_test.go b/pkg/kthena-router/router/router_test.go index ebe64444ea..acb8ed5ef4 100644 --- a/pkg/kthena-router/router/router_test.go +++ b/pkg/kthena-router/router/router_test.go @@ -3226,6 +3226,119 @@ func TestProxy_RetryBodyNotDrained(t *testing.T) { assert.Equal(t, http.StatusOK, accessCtx.UpstreamStatusCode) } +// TestProxy_RetryDropsFailedAttemptHeaders verifies that a backend which sends +// response headers and then fails before sending its declared body does not leak +// those headers into the retry. Headers staged before the attempt are restored, +// so the healthy response carries exactly one Content-Length and only its own +// Set-Cookie and X-Backend values. +func TestProxy_RetryDropsFailedAttemptHeaders(t *testing.T) { + var attempts int32 + backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/metrics" { + w.Header().Set("Content-Type", "text/plain; version=0.0.4") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, "") + return + } + if r.URL.Path == "/v1/models" { + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"data":[]}`) + return + } + if atomic.AddInt32(&attempts, 1) == 1 { + // Announce a 4096 byte body, then close the connection without it. + hijacker, ok := w.(http.Hijacker) + if !ok { + t.Error("backend does not support hijacking") + return + } + conn, buf, err := hijacker.Hijack() + if err != nil { + t.Errorf("hijack failed: %v", err) + return + } + defer conn.Close() + fmt.Fprint(buf, "HTTP/1.1 200 OK\r\nContent-Length: 4096\r\nContent-Type: application/json\r\nSet-Cookie: failed-backend=1\r\nX-Backend: failed\r\n\r\n") + if err := buf.Flush(); err != nil { + t.Errorf("flush failed: %v", err) + } + return + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Content-Length", "16") + w.Header().Set("X-Backend", "healthy") + w.WriteHeader(http.StatusOK) + fmt.Fprint(w, `{"id":"healthy"}`) + })) + defer backend.Close() + + backendURL, _ := url.Parse(backend.URL) + backendIP := backendURL.Hostname() + backendPort, _ := strconv.Atoi(backendURL.Port()) + + store := datastore.New() + router := NewRouter(store, "../scheduler/testdata/configmap.yaml", common.NewTransportRegistry()) + + modelServer := &aiv1alpha1.ModelServer{ + ObjectMeta: v1.ObjectMeta{Name: "ms-header-retry", Namespace: "default"}, + Spec: aiv1alpha1.ModelServerSpec{ + Model: func(s string) *string { return &s }("base-model"), + WorkloadPort: aiv1alpha1.WorkloadPort{Port: int32(backendPort)}, + InferenceEngine: "vLLM", + }, + } + pod1 := &corev1.Pod{ + ObjectMeta: v1.ObjectMeta{Name: "pod-header-retry-1", Namespace: "default"}, + Status: corev1.PodStatus{PodIP: backendIP, Phase: corev1.PodRunning}, + } + pod2 := &corev1.Pod{ + ObjectMeta: v1.ObjectMeta{Name: "pod-header-retry-2", Namespace: "default"}, + Status: corev1.PodStatus{PodIP: backendIP, Phase: corev1.PodRunning}, + } + modelRoute := &aiv1alpha1.ModelRoute{ + ObjectMeta: v1.ObjectMeta{Name: "mr-header-retry", Namespace: "default"}, + Spec: aiv1alpha1.ModelRouteSpec{ + ModelName: "header-retry-model", + Rules: []*aiv1alpha1.Rule{ + {TargetModels: []*aiv1alpha1.TargetModel{{ModelServerName: "ms-header-retry"}}}, + }, + }, + } + + store.AddOrUpdateModelServer(modelServer, sets.New( + types.NamespacedName{Name: "pod-header-retry-1", Namespace: "default"}, + types.NamespacedName{Name: "pod-header-retry-2", Namespace: "default"}, + )) + store.AddOrUpdatePod(pod1, []*aiv1alpha1.ModelServer{modelServer}) + store.AddOrUpdatePod(pod2, []*aiv1alpha1.ModelServer{modelServer}) + store.AddOrUpdateModelRoute(modelRoute) + + // Give pod-2 a waiting request so pod-1 scores higher and handles the first + // attempt, which is the one that fails after sending headers. + podInfo2 := store.GetPodInfo(types.NamespacedName{Name: "pod-header-retry-2", Namespace: "default"}) + require.NotNil(t, podInfo2) + podInfo2.RequestWaitingNum = 1 + + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request, _ = http.NewRequest("POST", "/v1/chat/completions", bytes.NewBufferString(`{"model": "header-retry-model", "prompt": "retry header test"}`)) + c.Request.Header.Set("Content-Type", "application/json") + // A header staged before forwarding stands in for router middleware. + c.Writer.Header().Set("X-Router-Middleware", "kept") + accessCtx := accesslog.NewAccessLogContext("header-retry-request", http.MethodPost, c.Request.URL.Path, c.Request.Proto, "header-retry-model") + c.Set(accesslog.AccessLogContextKey, accessCtx) + + router.HandlerFunc()(c) + + assert.Equal(t, http.StatusOK, w.Code) + assert.Equal(t, int32(2), atomic.LoadInt32(&attempts), "expected a failed attempt and a healthy retry") + assert.Equal(t, []string{"16"}, w.Header().Values("Content-Length"), "retry kept the failed attempt's Content-Length") + assert.Equal(t, []string{"healthy"}, w.Header().Values("X-Backend"), "retry kept the failed attempt's X-Backend") + assert.Empty(t, w.Header().Values("Set-Cookie"), "retry kept the failed attempt's cookies") + assert.Equal(t, []string{"kept"}, w.Header().Values("X-Router-Middleware"), "headers staged before forwarding were dropped") + assert.Equal(t, `{"id":"healthy"}`, w.Body.String()) +} + func TestRouter_HandlerFunc_ListModels(t *testing.T) { router, store, backend := setupTestRouter(t, nil) defer backend.Close()