From b53372e86a35d4ca3c1e2b7b99018654962209c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Gronowski?= <pawel.gronowski@docker.com>
Date: Thu, 1 Oct 2026 19:01:41 +0200
Subject: [PATCH] cli/command: Add WithContextResolver to override context
 selection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Callers that choose the context from configuration had to do it before
Initialize. That meant loading the config file and context store
themselves and handling --config a second time.

WithContextResolver runs a callback inside Initialize after the config
file and context store are loaded, but before telemetry and endpoint
initialization. A non-empty result replaces the selected context, an
empty result keeps normal selection, and an error aborts
initialization.

The API client and endpoint are resolved lazily, so they use the
overridden context.

Signed-off-by: Paweł Gronowski <pawel.gronowski@docker.com>
---
 cli/command/cli.go         | 13 +++++++++
 cli/command/cli_options.go | 13 +++++++++
 cli/command/cli_test.go    | 58 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+)

diff --git a/cli/command/cli.go b/cli/command/cli.go
index e405c2cc4f..0d2f4b506f 100644
--- a/cli/command/cli.go
+++ b/cli/command/cli.go
@@ -68,6 +68,7 @@ type DockerCli struct {
 	serverInfo         ServerInfo
 	contextStore       store.Store
 	currentContext     string
+	contextResolver    func(*DockerCli) (string, error)
 	init               sync.Once
 	initErr            error
 	dockerEndpoint     docker.Endpoint
@@ -208,6 +209,8 @@ func (cli *DockerCli) HooksEnabled() bool {
 
 // Initialize the dockerCli runs initialization that must happen after command
 // line flags are parsed.
+//
+//nolint:gocyclo
 func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions, ops ...CLIOption) error {
 	for _, o := range ops {
 		if err := o(cli); err != nil {
@@ -257,6 +260,16 @@ func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions, ops ...CLIOption)
 		},
 	}
 
+	if cli.contextResolver != nil {
+		contextName, err := cli.contextResolver(cli)
+		if err != nil {
+			return err
+		}
+		if contextName != "" {
+			cli.currentContext = contextName
+		}
+	}
+
 	// TODO(krissetto): pass ctx to the funcs instead of using this
 	if cli.enableGlobalMeter {
 		cli.createGlobalMeterProvider(cli.baseCtx)
diff --git a/cli/command/cli_options.go b/cli/command/cli_options.go
index 469932f9c3..733cfbdf64 100644
--- a/cli/command/cli_options.go
+++ b/cli/command/cli_options.go
@@ -84,6 +84,19 @@ func WithDefaultContextStoreConfig() CLIOption {
 	}
 }
 
+// WithContextResolver overrides context selection during [DockerCli.Initialize].
+// The resolver runs after configuration and the context store are loaded, but
+// before telemetry and endpoint initialization.
+// An empty result preserves normal context selection; an error aborts
+// initialization.
+// The resolver must not initialize the API client.
+func WithContextResolver(resolve func(*DockerCli) (string, error)) CLIOption {
+	return func(cli *DockerCli) error {
+		cli.contextResolver = resolve
+		return nil
+	}
+}
+
 // WithAPIClient configures the cli to use the given API client.
 func WithAPIClient(c client.APIClient) CLIOption {
 	return func(cli *DockerCli) error {
diff --git a/cli/command/cli_test.go b/cli/command/cli_test.go
index 27be8d348c..09c093d5dd 100644
--- a/cli/command/cli_test.go
+++ b/cli/command/cli_test.go
@@ -325,6 +325,64 @@ func TestInitializeShouldAlwaysCreateTheContextStore(t *testing.T) {
 	assert.Check(t, cli.ContextStore() != nil)
 }
 
+func TestInitializeContextResolver(t *testing.T) {
+	resolverErr := errors.New("resolver failed")
+
+	for _, tc := range []struct {
+		name    string
+		context string
+		err     error
+	}{
+		{name: "resolved context", context: "resolved"},
+		{name: "normal selection"},
+		{name: "resolver failure", err: resolverErr},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			originalConfigDir := config.Dir()
+			t.Cleanup(func() { config.SetDir(originalConfigDir) })
+			config.SetDir(t.TempDir())
+
+			configDir := t.TempDir()
+			cfg := configfile.New(filepath.Join(configDir, config.ConfigFileName))
+			cfg.Features = map[string]string{"cloud": "foobar"}
+			assert.NilError(t, cfg.Save())
+
+			cli, err := NewDockerCli()
+			assert.NilError(t, err)
+
+			var loadedConfig *configfile.ConfigFile
+			calls := 0
+
+			err = cli.Initialize(&flags.ClientOptions{ConfigDir: configDir, Context: "original"},
+				WithContextResolver(func(cli *DockerCli) (string, error) {
+					calls++
+					assert.Equal(t, config.Dir(), configDir)
+					loadedConfig = cli.ConfigFile()
+					assert.Equal(t, loadedConfig.Features["cloud"], "foobar")
+					assert.Assert(t, cli.ContextStore() != nil)
+					assert.Assert(t, cli.client == nil)
+					return tc.context, tc.err
+				}),
+			)
+			assert.Equal(t, calls, 1)
+			assert.Assert(t, cli.ConfigFile() == loadedConfig)
+			assert.Assert(t, cli.client == nil)
+
+			if tc.err != nil {
+				assert.ErrorIs(t, err, tc.err)
+				return
+			}
+			assert.NilError(t, err)
+
+			wantContext := tc.context
+			if wantContext == "" {
+				wantContext = "original"
+			}
+			assert.Equal(t, cli.CurrentContext(), wantContext)
+		})
+	}
+}
+
 func TestHooksEnabled(t *testing.T) {
 	t.Run("disabled by default", func(t *testing.T) {
 		// Make sure we don't depend on any existing ~/.docker/config.json