diff --git a/cli-plugins/plugin/plugin.go b/cli-plugins/plugin/plugin.go
index c04edbb549..5b1204f3cb 100644
--- a/cli-plugins/plugin/plugin.go
+++ b/cli-plugins/plugin/plugin.go
@@ -36,13 +36,7 @@ func RunPlugin(dockerCli *command.DockerCli, plugin *cobra.Command, meta manager
 	PersistentPreRunE = func(cmd *cobra.Command, _ []string) error {
 		var err error
 		persistentPreRunOnce.Do(func() {
-			cmdContext := cmd.Context()
-			// TODO: revisit and make sure this check makes sense
-			// see: https://github.com/docker/cli/pull/4599#discussion_r1422487271
-			if cmdContext == nil {
-				cmdContext = context.TODO()
-			}
-			ctx, cancel := context.WithCancel(cmdContext)
+			ctx, cancel := context.WithCancel(cmd.Context())
 			cmd.SetContext(ctx)
 			// Set up the context to cancel based on signalling via CLI socket.
 			socket.ConnectAndWait(cancel)
diff --git a/cmd/docker/builder_test.go b/cmd/docker/builder_test.go
index 20cd88af49..797d42140c 100644
--- a/cmd/docker/builder_test.go
+++ b/cmd/docker/builder_test.go
@@ -23,6 +23,9 @@ import (
 var pluginFilename = "docker-buildx"
 
 func TestBuildWithBuilder(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
 	testcases := []struct {
 		name         string
 		context      string
@@ -64,12 +67,16 @@ echo '{"SchemaVersion":"0.1.0","Vendor":"Docker Inc.","Version":"v0.6.3","ShortD
 	for _, tt := range testcases {
 		tt := tt
 		t.Run(tt.name, func(t *testing.T) {
+			ctx2, cancel2 := context.WithCancel(ctx)
+			defer cancel2()
+
 			if tt.builder != "" {
 				t.Setenv("BUILDX_BUILDER", tt.builder)
 			}
 
 			var b bytes.Buffer
 			dockerCli, err := command.NewDockerCli(
+				command.WithBaseContext(ctx2),
 				command.WithAPIClient(&fakeClient{}),
 				command.WithInputStream(discard),
 				command.WithCombinedStreams(&b),
@@ -126,6 +133,9 @@ func (c *fakeClient) Ping(_ context.Context) (types.Ping, error) {
 }
 
 func TestBuildkitDisabled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
 	t.Setenv("DOCKER_BUILDKIT", "0")
 
 	dir := fs.NewDir(t, t.Name(),
@@ -136,6 +146,7 @@ func TestBuildkitDisabled(t *testing.T) {
 	b := bytes.NewBuffer(nil)
 
 	dockerCli, err := command.NewDockerCli(
+		command.WithBaseContext(ctx),
 		command.WithAPIClient(&fakeClient{}),
 		command.WithInputStream(discard),
 		command.WithCombinedStreams(b),
@@ -163,6 +174,9 @@ func TestBuildkitDisabled(t *testing.T) {
 }
 
 func TestBuilderBroken(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
 	dir := fs.NewDir(t, t.Name(),
 		fs.WithFile(pluginFilename, `#!/bin/sh exit 1`, fs.WithMode(0o777)),
 	)
@@ -171,6 +185,7 @@ func TestBuilderBroken(t *testing.T) {
 	b := bytes.NewBuffer(nil)
 
 	dockerCli, err := command.NewDockerCli(
+		command.WithBaseContext(ctx),
 		command.WithAPIClient(&fakeClient{}),
 		command.WithInputStream(discard),
 		command.WithCombinedStreams(b),
@@ -199,6 +214,8 @@ func TestBuilderBroken(t *testing.T) {
 
 func TestBuilderBrokenEnforced(t *testing.T) {
 	t.Setenv("DOCKER_BUILDKIT", "1")
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
 
 	dir := fs.NewDir(t, t.Name(),
 		fs.WithFile(pluginFilename, `#!/bin/sh exit 1`, fs.WithMode(0o777)),
@@ -208,6 +225,7 @@ func TestBuilderBrokenEnforced(t *testing.T) {
 	b := bytes.NewBuffer(nil)
 
 	dockerCli, err := command.NewDockerCli(
+		command.WithBaseContext(ctx),
 		command.WithAPIClient(&fakeClient{}),
 		command.WithInputStream(discard),
 		command.WithCombinedStreams(b),
diff --git a/cmd/docker/docker.go b/cmd/docker/docker.go
index b87c24fba7..f68bda6fe7 100644
--- a/cmd/docker/docker.go
+++ b/cmd/docker/docker.go
@@ -29,6 +29,7 @@ import (
 
 func main() {
 	ctx := context.Background()
+
 	dockerCli, err := command.NewDockerCli(command.WithBaseContext(ctx))
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
@@ -352,7 +353,7 @@ func runDocker(ctx context.Context, dockerCli *command.DockerCli) error {
 	// We've parsed global args already, so reset args to those
 	// which remain.
 	cmd.SetArgs(args)
-	err = cmd.Execute()
+	err = cmd.ExecuteContext(ctx)
 
 	// If the command is being executed in an interactive terminal
 	// and hook are enabled, run the plugin hooks.
diff --git a/cmd/docker/docker_test.go b/cmd/docker/docker_test.go
index 084626252e..8fc7ade3a4 100644
--- a/cmd/docker/docker_test.go
+++ b/cmd/docker/docker_test.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"os"
 	"testing"
@@ -15,8 +16,10 @@ import (
 
 func TestClientDebugEnabled(t *testing.T) {
 	defer debug.Disable()
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
 
-	cli, err := command.NewDockerCli()
+	cli, err := command.NewDockerCli(command.WithBaseContext(ctx))
 	assert.NilError(t, err)
 	tcmd := newDockerCommand(cli)
 	tcmd.SetFlag("debug", "true")
@@ -39,7 +42,13 @@ func runCliCommand(t *testing.T, r io.ReadCloser, w io.Writer, args ...string) e
 	if w == nil {
 		w = io.Discard
 	}
-	cli, err := command.NewDockerCli(command.WithInputStream(r), command.WithCombinedStreams(w))
+	ctx, cancel := context.WithCancel(context.TODO())
+	defer cancel()
+
+	cli, err := command.NewDockerCli(
+		command.WithBaseContext(ctx),
+		command.WithInputStream(r),
+		command.WithCombinedStreams(w))
 	assert.NilError(t, err)
 	tcmd := newDockerCommand(cli)
 
diff --git a/docs/generate/generate.go b/docs/generate/generate.go
index a24aba1813..a12c7d4ccf 100644
--- a/docs/generate/generate.go
+++ b/docs/generate/generate.go
@@ -37,6 +37,7 @@ func gen(opts *options) error {
 		Use:   "docker [OPTIONS] COMMAND [ARG...]",
 		Short: "The base command for the Docker CLI.",
 	}
+
 	clientOpts, _ := cli.SetupRootCommand(cmd)
 	if err := dockerCLI.Initialize(clientOpts); err != nil {
 		return err
diff --git a/e2e/cli-plugins/plugins/nopersistentprerun/main.go b/e2e/cli-plugins/plugins/nopersistentprerun/main.go
index 5f07bf7e01..b90cc67aee 100644
--- a/e2e/cli-plugins/plugins/nopersistentprerun/main.go
+++ b/e2e/cli-plugins/plugins/nopersistentprerun/main.go
@@ -11,7 +11,7 @@ import (
 
 func main() {
 	plugin.Run(func(dockerCli command.Cli) *cobra.Command {
-		cmd := &cobra.Command{
+		return &cobra.Command{
 			Use:   "nopersistentprerun",
 			Short: "Testing without PersistentPreRun hooks",
 			// PersistentPreRunE: Not specified, we need to test that it works in the absence of an explicit call
@@ -25,7 +25,6 @@ func main() {
 				return nil
 			},
 		}
-		return cmd
 	},
 		manager.Metadata{
 			SchemaVersion: "0.1.0",