fix: enhance error handling for unsupported media types in pull operations

Author: ilopezluna

cmd/cli/commands/utils.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/docker/model-runner/cmd/cli/desktop"
 	"github.com/docker/model-runner/cmd/cli/pkg/standalone"
+	"github.com/docker/model-runner/pkg/distribution/distribution"
 	"github.com/docker/model-runner/pkg/distribution/oci/reference"
 	"github.com/docker/model-runner/pkg/inference/backends/vllm"
 	"github.com/moby/term"
@@ -53,7 +54,7 @@ func handleClientError(err error, message string) error {
 		var buf bytes.Buffer
 		printNextSteps(&buf, []string{enableVLLM})
 		return fmt.Errorf("%w\n%s", err, strings.TrimRight(buf.String(), "\n"))
-	} else if strings.Contains(err.Error(), "try upgrading") {
+	} else if errors.Is(err, distribution.ErrUnsupportedMediaType) {
 		// The model uses a newer config format than this client supports.
 		var buf bytes.Buffer
 		printNextSteps(&buf, []string{

cmd/cli/desktop/desktop.go

@@ -151,7 +151,16 @@ func (c *Client) Pull(model string, printer standalone.StatusPrinter) (string, b
 			} else {
 				bodyStr = strings.TrimSpace(string(body))
 			}
-			err := fmt.Errorf("pulling %s failed with status %s: %s", model, resp.Status, bodyStr)
+			var err error
+			if resp.StatusCode == http.StatusUnprocessableEntity {
+				// 422 means the model uses a config type this client does not
+				// support. Reattach the sentinel so callers can use errors.Is.
+				err = fmt.Errorf("pulling %s failed with status %s: %w: %s",
+					model, resp.Status, distribution.ErrUnsupportedMediaType, bodyStr)
+			} else {
+				err = fmt.Errorf("pulling %s failed with status %s: %s",
+					model, resp.Status, bodyStr)
+			}
 			// Only retry on gateway/proxy errors (502, 503, 504).
 			// Do not retry 500 (usually deterministic server errors) or
 			// 4xx (client errors including 422 for unsupported media type).
@@ -245,7 +254,7 @@ func (c *Client) withRetries(
 		}
 	}
 
-	return "", progressShown, fmt.Errorf("%w (failed after %d retries)", lastErr, maxRetries)
+	return "", progressShown, fmt.Errorf("%s failed after %d retries: %w", operationName, maxRetries, lastErr)
 }
 
 func (c *Client) Push(model string, printer standalone.StatusPrinter) (string, bool, error) {

cmd/cli/desktop/desktop_test.go

@@ -10,11 +10,18 @@ import (
 	"testing"
 
 	mockdesktop "github.com/docker/model-runner/cmd/cli/mocks"
+	"github.com/docker/model-runner/pkg/distribution/distribution"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
 )
 
+// errorReadCloser is an io.ReadCloser whose Read always returns an error.
+type errorReadCloser struct{ err error }
+
+func (e *errorReadCloser) Read(_ []byte) (int, error) { return 0, e.err }
+func (e *errorReadCloser) Close() error               { return nil }
+
 func TestPullRetryOnNetworkError(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
@@ -100,34 +107,52 @@ func TestPullNoRetryOn422Error(t *testing.T) {
 
 	printer := NewSimplePrinter(func(s string) {})
 	_, _, err := client.Pull(modelName, printer)
-	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "try upgrading")
+	require.Error(t, err)
+	// The sentinel must be preserved so callers can use errors.Is.
+	assert.True(t, errors.Is(err, distribution.ErrUnsupportedMediaType))
 }
 
-func TestPullRetryOn502Error(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	modelName := "test-model"
-	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
-	mockContext := NewContextForMock(mockClient)
-	client := New(mockContext)
-
-	// 502 Bad Gateway is a transient proxy error and should be retried.
-	gomock.InOrder(
-		mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
-			StatusCode: http.StatusBadGateway,
-			Body:       io.NopCloser(bytes.NewBufferString("Bad Gateway")),
-		}, nil),
-		mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
-			StatusCode: http.StatusOK,
-			Body:       io.NopCloser(bytes.NewBufferString(`{"type":"success","message":"Model pulled successfully"}`)),
-		}, nil),
-	)
+func TestPullRetriesOnTransientGatewayErrors(t *testing.T) {
+	// 502 and 504 are transient gateway/proxy errors and should be retried.
+	// Note: 503 is intercepted by doRequestWithAuthContext as ErrServiceUnavailable
+	// and is covered separately by TestPullRetryOnServiceUnavailable.
+	transientCodes := []struct {
+		code int
+		name string
+		body string
+	}{
+		{http.StatusBadGateway, "502 Bad Gateway", "Bad Gateway"},
+		{http.StatusGatewayTimeout, "504 Gateway Timeout", "Gateway Timeout"},
+	}
 
-	printer := NewSimplePrinter(func(s string) {})
-	_, _, err := client.Pull(modelName, printer)
-	assert.NoError(t, err)
+	for _, tc := range transientCodes {
+		t.Run(tc.name, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			defer ctrl.Finish()
+
+			mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+			mockContext := NewContextForMock(mockClient)
+			client := New(mockContext)
+
+			// First attempt fails with the transient error, second succeeds.
+			gomock.InOrder(
+				mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+					StatusCode: tc.code,
+					Body:       io.NopCloser(bytes.NewBufferString(tc.body)),
+				}, nil),
+				mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+					StatusCode: http.StatusOK,
+					Body: io.NopCloser(bytes.NewBufferString(
+						`{"type":"success","message":"Model pulled successfully"}`,
+					)),
+				}, nil),
+			)
+
+			printer := NewSimplePrinter(func(s string) {})
+			_, _, err := client.Pull("test-model", printer)
+			assert.NoError(t, err)
+		})
+	}
 }
 
 func TestPullRetryOnServiceUnavailable(t *testing.T) {
@@ -172,7 +197,7 @@ func TestPullMaxRetriesExhausted(t *testing.T) {
 	printer := NewSimplePrinter(func(s string) {})
 	_, _, err := client.Pull(modelName, printer)
 	assert.Error(t, err)
-	assert.Contains(t, err.Error(), "(failed after 3 retries)")
+	assert.Contains(t, err.Error(), "download failed after 3 retries")
 }
 
 func TestPushRetryOnNetworkError(t *testing.T) {
@@ -387,6 +412,27 @@ func TestIsTemplateIncompatibleError(t *testing.T) {
 	}
 }
 
+func TestPullBodyReadFailure(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// Response body read fails. Use a non-retryable 500 status so the test
+	// completes in a single attempt.
+	mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+		StatusCode: http.StatusInternalServerError,
+		Body:       &errorReadCloser{err: errors.New("connection reset")},
+	}, nil).Times(1)
+
+	printer := NewSimplePrinter(func(s string) {})
+	_, _, err := client.Pull("test-model", printer)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to read response body")
+}
+
 func TestDisplayProgressNonJSONLines(t *testing.T) {
 	// Simulate a proxy returning an HTML error page instead of a progress stream.
 	htmlBody := "<html><body><h1>502 Bad Gateway</h1></body></html>\n"

cmd/cli/desktop/progress.go

@@ -19,7 +19,9 @@ import (
 // DisplayProgress displays progress messages from a model pull/push operation
 // using Docker-style multi-line progress bars.
 // Returns the final message, whether progress was actually shown, and any error.
-func DisplayProgress(body io.Reader, printer standalone.StatusPrinter) (string, bool, error) {
+func DisplayProgress(
+	body io.Reader, printer standalone.StatusPrinter,
+) (finalMessage string, progressShown bool, retErr error) {
 	fd, isTerminal := printer.GetFdInfo()
 
 	// If not a terminal, fall back to simple line-by-line output
@@ -40,13 +42,22 @@ func DisplayProgress(body io.Reader, printer standalone.StatusPrinter) (string,
 		close(errCh)
 	}()
 
+	// Ensure the pipe is always closed and the display goroutine is always
+	// drained, even on early returns, to prevent goroutine leaks.
+	defer func() {
+		pw.Close()
+		if displayErr := <-errCh; retErr == nil &&
+			displayErr != nil && !errors.Is(displayErr, io.EOF) {
+			retErr = displayErr
+		}
+	}()
+
 	// Convert progress messages to JSONMessage format
 	scanner := bufio.NewScanner(body)
-	var finalMessage string
-	progressShown := false // Track if we actually showed any progress bars
 	// nonJSONBytes collects raw unparseable lines for error reporting,
 	// capped at maxNonJSONBytes to avoid large allocations.
 	var nonJSONBytes []byte
+	var nonJSONTruncated bool
 
 	for scanner.Scan() {
 		progressLine := scanner.Text()
@@ -58,15 +69,14 @@ func DisplayProgress(body io.Reader, printer standalone.StatusPrinter) (string,
 		if err := json.Unmarshal([]byte(html.UnescapeString(progressLine)), &progressMsg); err != nil {
 			// Collect unparseable lines (e.g. HTML error pages from proxies)
 			// so we can surface them if no valid progress arrives.
-			nonJSONBytes = appendNonJSONLine(nonJSONBytes, progressLine)
+			nonJSONBytes, nonJSONTruncated = appendNonJSONLine(nonJSONBytes, progressLine)
 			continue
 		}
 
 		switch progressMsg.Type {
 		case oci.TypeProgress:
 			progressShown = true // We're showing actual progress
 			if err := writeDockerProgress(pw, &progressMsg); err != nil {
-				pw.Close()
 				return "", false, err
 			}
 
@@ -80,33 +90,23 @@ func DisplayProgress(body io.Reader, printer standalone.StatusPrinter) (string,
 			printer.PrintErrf("Warning: %s\n", progressMsg.Message)
 
 		case oci.TypeError:
-			pw.Close()
 			return "", false, fmt.Errorf("%s", progressMsg.Message)
 		}
 	}
 
 	if err := scanner.Err(); err != nil {
-		pw.Close()
 		return "", false, err
 	}
 
 	// If we received only unparseable lines and no valid progress or success,
 	// surface the raw content as an error. This catches HTML error pages
 	// returned by proxies or CDNs in place of a proper progress stream.
 	if finalMessage == "" && !progressShown {
-		if err := unexpectedProgressDataError(nonJSONBytes); err != nil {
-			pw.Close()
+		if err := unexpectedProgressDataError(nonJSONBytes, nonJSONTruncated); err != nil {
 			return "", false, err
 		}
 	}
 
-	pw.Close()
-
-	// Wait for display to finish
-	if err := <-errCh; err != nil && !errors.Is(err, io.EOF) {
-		return finalMessage, progressShown, err
-	}
-
 	return finalMessage, progressShown, nil
 }
 
@@ -119,6 +119,7 @@ func displayProgressSimple(body io.Reader, printer standalone.StatusPrinter) (st
 	progressShown := false // Track if we actually showed any progress
 	// nonJSONBytes collects raw unparseable lines for error reporting.
 	var nonJSONBytes []byte
+	var nonJSONTruncated bool
 
 	for scanner.Scan() {
 		progressLine := scanner.Text()
@@ -129,7 +130,7 @@ func displayProgressSimple(body io.Reader, printer standalone.StatusPrinter) (st
 		var progressMsg oci.ProgressMessage
 		if err := json.Unmarshal([]byte(html.UnescapeString(progressLine)), &progressMsg); err != nil {
 			// Collect unparseable lines for error reporting.
-			nonJSONBytes = appendNonJSONLine(nonJSONBytes, progressLine)
+			nonJSONBytes, nonJSONTruncated = appendNonJSONLine(nonJSONBytes, progressLine)
 			continue
 		}
 
@@ -167,7 +168,7 @@ func displayProgressSimple(body io.Reader, printer standalone.StatusPrinter) (st
 
 	// Surface unparseable content if no valid progress was received.
 	if finalMessage == "" && !progressShown {
-		if err := unexpectedProgressDataError(nonJSONBytes); err != nil {
+		if err := unexpectedProgressDataError(nonJSONBytes, nonJSONTruncated); err != nil {
 			return "", false, err
 		}
 	}
@@ -289,29 +290,36 @@ func NewSimplePrinter(printFunc func(string)) standalone.StatusPrinter {
 const maxNonJSONBytes = 4096
 
 // appendNonJSONLine appends line (with a newline separator) to dst, enforcing
-// a hard cap of maxNonJSONBytes total to avoid large allocations.
-func appendNonJSONLine(dst []byte, line string) []byte {
+// a hard cap of maxNonJSONBytes total. Returns the updated slice and a boolean
+// indicating whether the line was truncated to fit within the cap.
+func appendNonJSONLine(dst []byte, line string) ([]byte, bool) {
 	if len(dst) >= maxNonJSONBytes {
-		return dst
+		return dst, true
 	}
 	if len(dst) > 0 {
 		dst = append(dst, '\n')
 	}
 	remaining := maxNonJSONBytes - len(dst)
-	if len(line) > remaining {
+	truncated := len(line) > remaining
+	if truncated {
 		line = line[:remaining]
 	}
-	return append(dst, line...)
+	return append(dst, line...), truncated
 }
 
 // unexpectedProgressDataError returns an error describing unexpected non-JSON
-// response data, or nil if nonJSONBytes is empty.
-func unexpectedProgressDataError(nonJSONBytes []byte) error {
+// response data, or nil if nonJSONBytes is empty. If truncated is true, a
+// marker is appended to indicate the response was cut off.
+func unexpectedProgressDataError(nonJSONBytes []byte, truncated bool) error {
 	if len(nonJSONBytes) == 0 {
 		return nil
 	}
+	msg := string(nonJSONBytes)
+	if truncated {
+		msg += "\n...[truncated]"
+	}
 	return fmt.Errorf(
 		"unexpected response from server (not valid progress data): %s",
-		string(nonJSONBytes),
+		msg,
 	)
 }