fix(upstream): retry a stream:false body cut after the 200 header

internal/upstream/loop_test.go

@@ -316,3 +316,52 @@ func TestStreamRelaysPersistent503(t *testing.T) {
 		t.Fatalf("body=%q", b)
 	}
 }
+
+// cutBody answers 200 and a promised body, then drops the connection half way: the client reads "unexpected EOF".
+func cutBody(t *testing.T, w http.ResponseWriter) {
+	t.Helper()
+	conn, buf, err := w.(http.Hijacker).Hijack()
+	if err != nil {
+		t.Fatal(err)
+	}
+	fmt.Fprint(buf, "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\n\r\ndata: {\"id\":\"c1\"")
+	_ = buf.Flush()
+	_ = conn.Close()
+}
+
+func TestRawRetriesABodyCutOffOnAnotherLane(t *testing.T) {
+	var n int64
+	c, _ := testClient(t, 2, func(w http.ResponseWriter, r *http.Request) {
+		if atomic.AddInt64(&n, 1) == 1 {
+			cutBody(t, w)
+			return
+		}
+		w.Header().Set("Content-Type", "text/event-stream")
+		fmt.Fprint(w, sseBody("whole"))
+	})
+	status, _, body := c.Raw("big-pickle", "id", []byte(`{"model":"big-pickle","messages":[]}`), nil)
+	if status != 200 || !strings.Contains(string(body), "whole") {
+		t.Fatalf("status=%d body=%s: a cut body must be fetched again, not returned", status, body)
+	}
+	if atomic.LoadInt64(&n) != 2 {
+		t.Fatalf("attempts=%d, want 2", n)
+	}
+}
+
+func TestRawNeverReportsACutBodyAsSuccess(t *testing.T) {
+	var n int64
+	c, _ := testClient(t, 2, func(w http.ResponseWriter, r *http.Request) {
+		atomic.AddInt64(&n, 1)
+		cutBody(t, w)
+	})
+	status, _, body := c.Raw("big-pickle", "id", []byte(`{"model":"big-pickle","messages":[]}`), nil)
+	if status != 502 {
+		t.Fatalf("status=%d body=%s: an answer that never arrived whole is a 502, not a 200", status, body)
+	}
+	if !strings.Contains(string(body), `"error":"`) || len(body) < 20 {
+		t.Fatalf("body=%s should carry the read error", body)
+	}
+	if got := atomic.LoadInt64(&n); got < 2 || got > 4 {
+		t.Fatalf("attempts=%d: retried on other lanes, but bounded (2 passes over 2 lanes)", got)
+	}
+}

internal/upstream/upstream.go

@@ -334,9 +334,14 @@ func (c *Client) Raw(model, identity string, fwd []byte, in http.Header) (int, h
 			outH = hdr
 			b, derr := drain(resp)
 			outB = b
-			return !retryable(st), derr
+			// a body cut mid-way is not an answer: nothing went downstream yet, so the loop may try another lane
+			return derr == nil && !retryable(st), derr
 		})
 	if err != nil {
+		if status < 300 {
+			// the header said 200 but the body never arrived whole: the caller must see a failure it can retry
+			status = 502
+		}
 		return status, jsonHeader(), []byte(`{"error":` + quote(err.Error()) + `}`)
 	}
 	if outH == nil {
@@ -365,7 +370,7 @@ func (c *Client) GetURL(rawURL string, budget time.Duration) (int, []byte) {
 		func(resp *http.Response, st int, hdr http.Header) (bool, error) {
 			b, derr := drain(resp)
 			out = b
-			return !retryable(st), derr
+			return derr == nil && !retryable(st), derr
 		})
 	return status, out
 }
@@ -624,6 +629,16 @@ func (c *Client) loopBudget(model, identity, method, path string, fwd []byte, in
 				continue
 			}
 			return st, cerr
+		case cerr != nil && !stop && st < 300:
+			// The header was fine but the body broke off ("unexpected EOF") and nothing was committed downstream:
+			// a far-end failure like a 502, worth another exit, bounded like the 5xx case.
+			lastErr = cerr.Error()
+			c.Lanes.Emit(lanes.Proof{T: time.Now(), Lane: lane.Index, Country: lane.Country, IP: lane.ExitIP, Model: model, Status: 502, Note: "body cut: " + lastErr})
+			farFails++
+			if farFails >= 2*len(c.Lanes.Lanes()) || time.Now().After(deadline) {
+				return 502, cerr
+			}
+			continue
 		default:
 			// committed answer: score the exit and leave a proof row
 			c.Lanes.NoteResult(lane.Country, st)