Automatic retries

When you run cmd/cli/model-cli pull gemma3 and the download gets interrupted:

1. The CLI detects the interruption (network error, timeout, etc.)
2. It waits 1 second and automatically retries
3. If it fails again, it waits 2 seconds and retries
4. If it fails a third time, it waits 4 seconds and retries
5. After 3 failed retries, it reports the error to the user

The user will see messages like:
Retrying download (attempt 1/3) in 1s...
Retrying download (attempt 2/3) in 2s...

Signed-off-by: Eric Curtin <eric.curtin@docker.com>

Author: ericcurtin

cmd/cli/desktop/desktop.go

@@ -106,61 +106,160 @@ func (c *Client) Status() Status {
 
 func (c *Client) Pull(model string, ignoreRuntimeMemoryCheck bool, printer standalone.StatusPrinter) (string, bool, error) {
 	model = normalizeHuggingFaceModelName(model)
-	jsonData, err := json.Marshal(dmrm.ModelCreateRequest{From: model, IgnoreRuntimeMemoryCheck: ignoreRuntimeMemoryCheck})
-	if err != nil {
-		return "", false, fmt.Errorf("error marshaling request: %w", err)
+
+	return c.withRetries("download", 3, printer, func(attempt int) (string, bool, error, bool) {
+		jsonData, err := json.Marshal(dmrm.ModelCreateRequest{From: model, IgnoreRuntimeMemoryCheck: ignoreRuntimeMemoryCheck})
+		if err != nil {
+			// Marshaling errors are not retryable
+			return "", false, fmt.Errorf("error marshaling request: %w", err), false
+		}
+
+		createPath := inference.ModelsPrefix + "/create"
+		resp, err := c.doRequest(
+			http.MethodPost,
+			createPath,
+			bytes.NewReader(jsonData),
+		)
+		if err != nil {
+			// Only retry on network errors, not on client errors
+			if isRetryableError(err) {
+				return "", false, c.handleQueryError(err, createPath), true
+			}
+			return "", false, c.handleQueryError(err, createPath), false
+		}
+		// Close response body explicitly at the end of this attempt, not deferred
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			body, _ := io.ReadAll(resp.Body)
+			err := fmt.Errorf("pulling %s failed with status %s: %s", model, resp.Status, string(body))
+			// Only retry on server errors (5xx), not client errors (4xx)
+			shouldRetry := resp.StatusCode >= 500 && resp.StatusCode < 600
+			return "", false, err, shouldRetry
+		}
+
+		// Use Docker-style progress display
+		message, shown, err := DisplayProgress(resp.Body, printer)
+		if err != nil {
+			// Retry on progress display errors (likely network interruption)
+			shouldRetry := isRetryableError(err)
+			return "", shown, err, shouldRetry
+		}
+
+		return message, shown, nil, false
+	})
+}
+
+// isRetryableError determines if an error is retryable (network-related)
+func isRetryableError(err error) bool {
+	if err == nil {
+		return false
 	}
 
-	createPath := inference.ModelsPrefix + "/create"
-	resp, err := c.doRequest(
-		http.MethodPost,
-		createPath,
-		bytes.NewReader(jsonData),
-	)
-	if err != nil {
-		return "", false, c.handleQueryError(err, createPath)
+	// First check for specific error types using errors.Is
+	if errors.Is(err, context.DeadlineExceeded) ||
+		errors.Is(err, io.ErrUnexpectedEOF) ||
+		errors.Is(err, io.EOF) ||
+		errors.Is(err, ErrServiceUnavailable) {
+		return true
 	}
-	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusOK {
-		body, _ := io.ReadAll(resp.Body)
-		return "", false, fmt.Errorf("pulling %s failed with status %s: %s", model, resp.Status, string(body))
+	// Fall back to string matching for network errors that don't have specific types
+	// This is necessary because many network errors are only available as strings
+	errStr := err.Error()
+	retryablePatterns := []string{
+		"connection refused",
+		"connection reset",
+		"broken pipe",
+		"timeout",
+		"temporary failure",
+		"no such host",
+		"no route to host",
+		"network is unreachable",
+		"i/o timeout",
 	}
 
-	// Use Docker-style progress display
-	message, progressShown, err := DisplayProgress(resp.Body, printer)
-	if err != nil {
-		return "", progressShown, err
+	for _, pattern := range retryablePatterns {
+		if strings.Contains(strings.ToLower(errStr), pattern) {
+			return true
+		}
 	}
 
-	return message, progressShown, nil
+	return false
+}
+
+// withRetries executes an operation with automatic retry logic for transient failures
+func (c *Client) withRetries(
+	operationName string,
+	maxRetries int,
+	printer standalone.StatusPrinter,
+	operation func(attempt int) (message string, progressShown bool, err error, shouldRetry bool),
+) (string, bool, error) {
+	var lastErr error
+	var progressShown bool
+
+	for attempt := 0; attempt <= maxRetries; attempt++ {
+		if attempt > 0 {
+			// Calculate exponential backoff: 2^(attempt-1) seconds (1s, 2s, 4s)
+			backoffDuration := time.Duration(1<<uint(attempt-1)) * time.Second
+			printer.PrintErrf("Retrying %s (attempt %d/%d) in %v...\n", operationName, attempt, maxRetries, backoffDuration)
+			time.Sleep(backoffDuration)
+		}
+
+		message, shown, err, shouldRetry := operation(attempt)
+		progressShown = progressShown || shown
+
+		if err == nil {
+			return message, progressShown, nil
+		}
+
+		lastErr = err
+		if !shouldRetry {
+			return "", progressShown, err
+		}
+	}
+
+	return "", progressShown, fmt.Errorf("failed to %s after %d retries: %w", operationName, maxRetries, lastErr)
 }
 
 func (c *Client) Push(model string, printer standalone.StatusPrinter) (string, bool, error) {
 	model = normalizeHuggingFaceModelName(model)
-	pushPath := inference.ModelsPrefix + "/" + model + "/push"
-	resp, err := c.doRequest(
-		http.MethodPost,
-		pushPath,
-		nil, // Assuming no body is needed for the push request
-	)
-	if err != nil {
-		return "", false, c.handleQueryError(err, pushPath)
-	}
-	defer resp.Body.Close()
 
-	if resp.StatusCode != http.StatusOK {
-		body, _ := io.ReadAll(resp.Body)
-		return "", false, fmt.Errorf("pushing %s failed with status %s: %s", model, resp.Status, string(body))
-	}
+	return c.withRetries("push", 3, printer, func(attempt int) (string, bool, error, bool) {
+		pushPath := inference.ModelsPrefix + "/" + model + "/push"
+		resp, err := c.doRequest(
+			http.MethodPost,
+			pushPath,
+			nil, // Assuming no body is needed for the push request
+		)
+		if err != nil {
+			// Only retry on network errors, not on client errors
+			if isRetryableError(err) {
+				return "", false, c.handleQueryError(err, pushPath), true
+			}
+			return "", false, c.handleQueryError(err, pushPath), false
+		}
+		// Close response body explicitly at the end of this attempt, not deferred
+		defer resp.Body.Close()
 
-	// Use Docker-style progress display
-	message, progressShown, err := DisplayProgress(resp.Body, printer)
-	if err != nil {
-		return "", progressShown, err
-	}
+		if resp.StatusCode != http.StatusOK {
+			body, _ := io.ReadAll(resp.Body)
+			err := fmt.Errorf("pushing %s failed with status %s: %s", model, resp.Status, string(body))
+			// Only retry on server errors (5xx), not client errors (4xx)
+			shouldRetry := resp.StatusCode >= 500 && resp.StatusCode < 600
+			return "", false, err, shouldRetry
+		}
+
+		// Use Docker-style progress display
+		message, shown, err := DisplayProgress(resp.Body, printer)
+		if err != nil {
+			// Retry on progress display errors (likely network interruption)
+			shouldRetry := isRetryableError(err)
+			return "", shown, err, shouldRetry
+		}
 
-	return message, progressShown, nil
+		return message, shown, nil, false
+	})
 }
 
 func (c *Client) List() ([]dmrm.Model, error) {

cmd/cli/desktop/desktop_test.go

@@ -2,7 +2,9 @@ package desktop
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
+	"errors"
 	"io"
 	"net/http"
 	"testing"
@@ -228,3 +230,171 @@ func TestInspectOpenAIHuggingFaceModel(t *testing.T) {
 	assert.Equal(t, expectedLowercase, model.ID)
 }
 
+func TestPullRetryOnNetworkError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// First two attempts fail with network error, third succeeds
+	gomock.InOrder(
+		mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.ErrUnexpectedEOF),
+		mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.ErrUnexpectedEOF),
+		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(modelName, false, printer)
+	assert.NoError(t, err)
+}
+
+func TestPullNoRetryOn4xxError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// Should not retry on 404 (client error)
+	mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+		StatusCode: http.StatusNotFound,
+		Body:       io.NopCloser(bytes.NewBufferString("Model not found")),
+	}, nil).Times(1)
+
+	printer := NewSimplePrinter(func(s string) {})
+	_, _, err := client.Pull(modelName, false, printer)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "Model not found")
+}
+
+func TestPullRetryOn5xxError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// First attempt fails with 500, second succeeds
+	gomock.InOrder(
+		mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+			StatusCode: http.StatusInternalServerError,
+			Body:       io.NopCloser(bytes.NewBufferString("Internal server error")),
+		}, 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(modelName, false, printer)
+	assert.NoError(t, err)
+}
+
+func TestPullRetryOnServiceUnavailable(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// First attempt fails with 503 (converted to ErrServiceUnavailable), second succeeds
+	// Note: 503 is handled specially in doRequestWithAuthContext and returns ErrServiceUnavailable
+	gomock.InOrder(
+		mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+			StatusCode: http.StatusServiceUnavailable,
+			Body:       io.NopCloser(bytes.NewBufferString("Service temporarily unavailable")),
+		}, 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(modelName, false, printer)
+	assert.NoError(t, err)
+}
+
+func TestPullMaxRetriesExhausted(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// All 4 attempts (1 initial + 3 retries) fail with network error
+	mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.EOF).Times(4)
+
+	printer := NewSimplePrinter(func(s string) {})
+	_, _, err := client.Pull(modelName, false, printer)
+	assert.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to download after 3 retries")
+}
+
+func TestPushRetryOnNetworkError(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	modelName := "test-model"
+	mockClient := mockdesktop.NewMockDockerHttpClient(ctrl)
+	mockContext := NewContextForMock(mockClient)
+	client := New(mockContext)
+
+	// First attempt fails with network error, second succeeds
+	gomock.InOrder(
+		mockClient.EXPECT().Do(gomock.Any()).Return(nil, io.ErrUnexpectedEOF),
+		mockClient.EXPECT().Do(gomock.Any()).Return(&http.Response{
+			StatusCode: http.StatusOK,
+			Body:       io.NopCloser(bytes.NewBufferString(`{"type":"success","message":"Model pushed successfully"}`)),
+		}, nil),
+	)
+
+	printer := NewSimplePrinter(func(s string) {})
+	_, _, err := client.Push(modelName, printer)
+	assert.NoError(t, err)
+}
+
+func TestIsRetryableError(t *testing.T) {
+	tests := []struct {
+		name     string
+		err      error
+		expected bool
+	}{
+		{"nil error", nil, false},
+		{"EOF error", io.EOF, true},
+		{"UnexpectedEOF error", io.ErrUnexpectedEOF, true},
+		{"connection reset in string", errors.New("some error: connection reset by peer"), true},
+		{"timeout in string", errors.New("operation failed: i/o timeout"), true},
+		{"connection refused", errors.New("dial tcp: connection refused"), true},
+		{"broken pipe", errors.New("write: broken pipe"), true},
+		{"network unreachable", errors.New("network is unreachable"), true},
+		{"no such host", errors.New("lookup failed: no such host"), true},
+		{"no route to host", errors.New("read tcp: no route to host"), true},
+		{"generic non-retryable error", errors.New("a generic non-retryable error"), false},
+		{"service unavailable error", ErrServiceUnavailable, true},
+		{"deadline exceeded", context.DeadlineExceeded, true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isRetryableError(tt.err)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+