fix: validate context size range and update backend configuration handling

Author: ilopezluna

cmd/cli/commands/compose.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"net"
 	"slices"
 	"strconv"
@@ -85,14 +86,20 @@ func newUpCommand() *cobra.Command {
 			}
 
 			for _, model := range models {
-				size := int32(ctxSize)
+				backendConfig := inference.BackendConfiguration{
+					Speculative: speculativeConfig,
+				}
+				if cmd.Flags().Changed("context-size") {
+					if ctxSize > math.MaxInt32 || ctxSize < math.MinInt32 {
+						return fmt.Errorf("context-size %d is out of range (must be between %d and %d)", ctxSize, math.MinInt32, math.MaxInt32)
+					}
+					size := int32(ctxSize)
+					backendConfig.ContextSize = &size
+				}
 				if err := desktopClient.ConfigureBackend(scheduling.ConfigureRequest{
-					Model: model,
-					BackendConfiguration: inference.BackendConfiguration{
-						ContextSize: &size,
-						Speculative: speculativeConfig,
-					},
-					RawRuntimeFlags: rawRuntimeFlags,
+					Model:                model,
+					BackendConfiguration: backendConfig,
+					RawRuntimeFlags:      rawRuntimeFlags,
 				}); err != nil {
 					configErrFmtString := "failed to configure backend for model %s with context-size %d  and runtime-flags %s"
 					_ = sendErrorf(configErrFmtString+": %v", model, ctxSize, rawRuntimeFlags, err)

cmd/cli/commands/compose_test.go

@@ -1,13 +1,104 @@
 package commands
 
 import (
+	"fmt"
+	"math"
 	"testing"
 
 	"github.com/docker/model-runner/pkg/inference"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
+// TestUpCommandContextSizeFlagBehavior verifies that the --context-size flag on
+// the compose up command is not "changed" by default (i.e. nil ContextSize
+// should be sent when the flag is absent) and is marked as changed after an
+// explicit value is provided.
+func TestUpCommandContextSizeFlagBehavior(t *testing.T) {
+	t.Run("context-size flag not changed by default", func(t *testing.T) {
+		cmd := newUpCommand()
+		// Parse with just the required --model flag — no --context-size.
+		err := cmd.ParseFlags([]string{"--model", "mymodel"})
+		require.NoError(t, err)
+		// The flag must NOT be marked as changed so that ContextSize is omitted
+		// from the configure request (i.e. remains nil).
+		assert.False(t, cmd.Flags().Changed("context-size"),
+			"context-size must not be Changed when the flag is absent")
+	})
+
+	t.Run("context-size flag changed after explicit value", func(t *testing.T) {
+		cmd := newUpCommand()
+		err := cmd.ParseFlags([]string{"--model", "mymodel", "--context-size", "4096"})
+		require.NoError(t, err)
+		assert.True(t, cmd.Flags().Changed("context-size"),
+			"context-size must be Changed when explicitly provided")
+	})
+
+	t.Run("context-size flag changed with unlimited value -1", func(t *testing.T) {
+		cmd := newUpCommand()
+		err := cmd.ParseFlags([]string{"--model", "mymodel", "--context-size", "-1"})
+		require.NoError(t, err)
+		assert.True(t, cmd.Flags().Changed("context-size"),
+			"context-size must be Changed when explicitly set to -1 (unlimited)")
+	})
+
+	t.Run("ContextSize is nil in BackendConfiguration when flag not set", func(t *testing.T) {
+		cmd := newUpCommand()
+		require.NoError(t, cmd.ParseFlags([]string{"--model", "mymodel"}))
+		// Simulate the logic in compose.go RunE: only add ContextSize when Changed.
+		backendConfig := inference.BackendConfiguration{}
+		if cmd.Flags().Changed("context-size") {
+			size := int32(-1) // default value
+			backendConfig.ContextSize = &size
+		}
+		assert.Nil(t, backendConfig.ContextSize,
+			"ContextSize must be nil in BackendConfiguration when --context-size is not provided")
+	})
+
+	t.Run("ContextSize is non-nil in BackendConfiguration when flag is set", func(t *testing.T) {
+		cmd := newUpCommand()
+		require.NoError(t, cmd.ParseFlags([]string{"--model", "mymodel", "--context-size", "64000"}))
+		ctxSize, err := cmd.Flags().GetInt64("context-size")
+		require.NoError(t, err)
+		backendConfig := inference.BackendConfiguration{}
+		if cmd.Flags().Changed("context-size") {
+			size := int32(ctxSize)
+			backendConfig.ContextSize = &size
+		}
+		require.NotNil(t, backendConfig.ContextSize,
+			"ContextSize must be non-nil when --context-size is provided")
+		assert.Equal(t, int32(64000), *backendConfig.ContextSize)
+	})
+
+	t.Run("context-size above int32 max is out of range", func(t *testing.T) {
+		tooBig := int64(math.MaxInt32) + 1
+		cmd := newUpCommand()
+		require.NoError(t, cmd.ParseFlags([]string{"--model", "mymodel", "--context-size", fmt.Sprintf("%d", tooBig)}))
+		ctxSize, err := cmd.Flags().GetInt64("context-size")
+		require.NoError(t, err)
+		require.True(t, cmd.Flags().Changed("context-size"))
+		// Simulate the range check from compose.go RunE.
+		if ctxSize > math.MaxInt32 || ctxSize < math.MinInt32 {
+			// Expected: would return an error in RunE.
+			return
+		}
+		t.Fatal("expected out-of-range check to trigger for value above MaxInt32")
+	})
+
+	t.Run("context-size below int32 min is out of range", func(t *testing.T) {
+		tooSmall := int64(math.MinInt32) - 1
+		cmd := newUpCommand()
+		require.NoError(t, cmd.ParseFlags([]string{"--model", "mymodel", "--context-size", fmt.Sprintf("%d", tooSmall)}))
+		ctxSize, err := cmd.Flags().GetInt64("context-size")
+		require.NoError(t, err)
+		require.True(t, cmd.Flags().Changed("context-size"))
+		if ctxSize > math.MaxInt32 || ctxSize < math.MinInt32 {
+			return
+		}
+		t.Fatal("expected out-of-range check to trigger for value below MinInt32")
+	})
+}
+
 func TestParseBackendMode(t *testing.T) {
 	tests := []struct {
 		name        string

pkg/inference/backends/sglang/sglang_config.go

@@ -72,6 +72,9 @@ func GetContextLength(modelCfg types.ModelConfig, backendCfg *inference.BackendC
 		return backendCfg.ContextSize
 	}
 	// Fallback to model config (set at packaging time via docker model package --context-size)
+	if modelCfg == nil {
+		return nil
+	}
 	if cs := modelCfg.GetContextSize(); cs != nil && *cs > 0 {
 		return cs
 	}

pkg/inference/backends/sglang/sglang_config_test.go

@@ -265,6 +265,18 @@ func TestGetContextLength(t *testing.T) {
 			},
 			expectedValue: nil,
 		},
+		{
+			name:          "nil model config with backend config",
+			modelCfg:      nil,
+			backendCfg:    &inference.BackendConfiguration{ContextSize: int32ptr(4096)},
+			expectedValue: int32ptr(4096),
+		},
+		{
+			name:          "nil model config without backend config",
+			modelCfg:      nil,
+			backendCfg:    nil,
+			expectedValue: nil,
+		},
 	}
 
 	for _, tt := range tests {