From cd98ded6ea2b821c8b14c53372e73a8f3f42b1d2 Mon Sep 17 00:00:00 2001
From: bwplotka <bwplotka@gmail.com>
Date: Fri, 12 Dec 2025 22:43:51 +0000
Subject: [PATCH] test: add regression test against remote write handler bad
 response stats

Signed-off-by: bwplotka <bwplotka@gmail.com>
---
 storage/remote/client.go             |  3 ++
 storage/remote/write_handler.go      |  4 ++
 storage/remote/write_handler_test.go | 72 ++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/storage/remote/client.go b/storage/remote/client.go
index c535ea3425..0f2b5ddca6 100644
--- a/storage/remote/client.go
+++ b/storage/remote/client.go
@@ -301,6 +301,9 @@ func (c *Client) Store(ctx context.Context, req []byte, attempt int) (WriteRespo
 		_ = httpResp.Body.Close()
 	}()
 
+	// NOTE(bwplotka): Only PRW2 spec defines response HTTP headers. However, spec does not block
+	// PRW1 from sending them too for reliability. Support this case.
+	//
 	// TODO(bwplotka): Pass logger and emit debug on error?
 	// Parsing error means there were some response header values we can't parse,
 	// we can continue handling.
diff --git a/storage/remote/write_handler.go b/storage/remote/write_handler.go
index f8296b4a80..2bc65e8286 100644
--- a/storage/remote/write_handler.go
+++ b/storage/remote/write_handler.go
@@ -96,6 +96,10 @@ func isHistogramValidationError(err error) bool {
 }
 
 // Store implements remoteapi.writeStorage interface.
+// TODO(bwplotka): Improve remoteapi.Store API. Right now it's confusing if PRWv1 flows should use WriteResponse or not.
+// If it's not filled, it will be "confirmed zero" which caused partial error reporting on client side in the past.
+// Temporary fix was done to only care about WriteResponse stats for PRW2 (see https://github.com/prometheus/client_golang/pull/1927
+// but better approach would be to only confirm if explicit stats were injected.
 func (h *writeHandler) Store(r *http.Request, msgType remoteapi.WriteMessageType) (*remoteapi.WriteResponse, error) {
 	// Store receives request with decompressed content in body.
 	body, err := io.ReadAll(r.Body)
diff --git a/storage/remote/write_handler_test.go b/storage/remote/write_handler_test.go
index f1c064c64d..2610142db9 100644
--- a/storage/remote/write_handler_test.go
+++ b/storage/remote/write_handler_test.go
@@ -1510,3 +1510,75 @@ func TestHistogramsReduction(t *testing.T) {
 		})
 	}
 }
+
+// Regression test for https://github.com/prometheus/prometheus/issues/17659
+func TestRemoteWriteHandler_ResponseStats(t *testing.T) {
+	payloadV1, _, _, err := buildWriteRequest(nil, writeRequestFixture.Timeseries, nil, nil, nil, nil, "snappy")
+	require.NoError(t, err)
+	payloadV2, _, _, err := buildV2WriteRequest(nil, writeV2RequestFixture.Timeseries, writeV2RequestFixture.Symbols, nil, nil, nil, "snappy")
+	require.NoError(t, err)
+
+	for _, tt := range []struct {
+		msgType            remoteapi.WriteMessageType
+		forceInjectHeaders bool
+		expectHeaders      bool
+	}{
+		{
+			msgType: remoteapi.WriteV1MessageType,
+		},
+		{
+			msgType:            remoteapi.WriteV1MessageType,
+			forceInjectHeaders: true,
+			expectHeaders:      true,
+		},
+		{
+			msgType:       remoteapi.WriteV2MessageType,
+			expectHeaders: true,
+		},
+	} {
+		t.Run(fmt.Sprintf("msg=%v/force-inject-headers=%v", tt.msgType, tt.forceInjectHeaders), func(t *testing.T) {
+			// Setup server side.
+			appendable := &mockAppendable{}
+			handler := NewWriteHandler(
+				promslog.NewNopLogger(),
+				nil,
+				appendable,
+				[]remoteapi.WriteMessageType{remoteapi.WriteV1MessageType, remoteapi.WriteV2MessageType},
+				false,
+				false,
+				false,
+			)
+
+			if tt.forceInjectHeaders {
+				base := handler
+				handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					base.ServeHTTP(w, r)
+
+					// Inject response header. This simulates PRWv1 server that uses PRWv2 response headers
+					// for confirmation of samples. This is not against spec and we support it.
+					w.Header().Set(rw20WrittenSamplesHeader, fmt.Sprintf("%d", len(appendable.samples)))
+				})
+			}
+
+			srv := httptest.NewServer(handler)
+
+			// Send message and do the parse response flow.
+			c := &Client{Client: srv.Client(), urlString: srv.URL, timeout: 5 * time.Minute, writeProtoMsg: tt.msgType}
+
+			payload := payloadV2
+			if tt.msgType == remoteapi.WriteV1MessageType {
+				payload = payloadV1
+			}
+			stats, err := c.Store(t.Context(), payload, 0)
+			require.NoError(t, err)
+
+			fmt.Println(stats)
+			if tt.expectHeaders {
+				require.True(t, stats.Confirmed)
+				require.Equal(t, len(appendable.samples), stats.Samples)
+			} else {
+				require.False(t, stats.Confirmed)
+			}
+		})
+	}
+}