Always read request body to facilitate retries

pkg/proxy/handler_test.go

@@ -19,13 +19,11 @@ import (
 func TestHandler(t *testing.T) {
 	const (
 		model1 = "model1"
-		model2 = "model2"
 
 		maxRetries = 3
 	)
 	models := map[string]string{
 		model1: "deploy1",
-		model2: "deploy2",
 	}
 
 	specs := map[string]struct {

pkg/proxy/request.go

@@ -60,17 +60,20 @@ func (p *proxyRequest) done() {
 // attempts to unmarshal the request body as JSON and extract the
 // .model field.
 func (pr *proxyRequest) parseModel() error {
-	pr.model = pr.r.Header.Get("X-Model")
-	if pr.model != "" {
-		return nil
-	}
-
+	// The request body always needs to be read (even when the model is set in the
+	// header) to facilitate retries. Otherwise the first request that is proxied
+	// will read the body and it will be empty for the subsequent proxy requests.
 	var err error
 	pr.body, err = io.ReadAll(pr.r.Body)
 	if err != nil {
 		return fmt.Errorf("read: %w", err)
 	}
 
+	pr.model = pr.r.Header.Get("X-Model")
+	if pr.model != "" {
+		return nil
+	}
+
 	var payload struct {
 		Model string `json:"model"`
 	}
@@ -116,11 +119,9 @@ func (pr *proxyRequest) setStatus(w http.ResponseWriter, code int) {
 
 // httpRequest returns a new http.Request that is a clone of the original
 // request, preserving the original request body even if it was already
-// read (i.e. if the body was inspected to determine the model).
+// read.
 func (pr *proxyRequest) httpRequest() *http.Request {
 	clone := pr.r.Clone(pr.r.Context())
-	if pr.body != nil {
-		clone.Body = io.NopCloser(bytes.NewReader(pr.body))
-	}
+	clone.Body = io.NopCloser(bytes.NewReader(pr.body))
 	return clone
 }