diff --git a/commands/bake.go b/commands/bake.go
index 5cfa397991f4..e17a8ebc7dad 100644
--- a/commands/bake.go
+++ b/commands/bake.go
@@ -24,7 +24,6 @@ import (
 	"github.com/docker/buildx/util/tracing"
 	"github.com/docker/cli/cli/command"
 	"github.com/moby/buildkit/identity"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/moby/buildkit/util/progress/progressui"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
@@ -43,9 +42,7 @@ type bakeOptions struct {
 	exportLoad   bool
 }
 
-func runBake(dockerCli command.Cli, targets []string, in bakeOptions, cFlags commonFlags) (err error) {
-	ctx := appcontext.Context()
-
+func runBake(ctx context.Context, dockerCli command.Cli, targets []string, in bakeOptions, cFlags commonFlags) (err error) {
 	mp, report, err := metrics.MeterProvider(dockerCli)
 	if err != nil {
 		return err
@@ -275,7 +272,7 @@ func bakeCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 			options.builder = rootOpts.builder
 			options.metadataFile = cFlags.metadataFile
 			// Other common flags (noCache, pull and progress) are processed in runBake function.
-			return runBake(dockerCli, args, options, cFlags)
+			return runBake(cmd.Context(), dockerCli, args, options, cFlags)
 		},
 		ValidArgsFunction: completion.BakeTargets(options.files),
 	}
diff --git a/commands/build.go b/commands/build.go
index a13446845b74..e142befd988b 100644
--- a/commands/build.go
+++ b/commands/build.go
@@ -46,7 +46,6 @@ import (
 	"github.com/moby/buildkit/frontend/subrequests/outline"
 	"github.com/moby/buildkit/frontend/subrequests/targets"
 	"github.com/moby/buildkit/solver/errdefs"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/moby/buildkit/util/grpcerrors"
 	"github.com/moby/buildkit/util/progress/progressui"
 	"github.com/morikuni/aec"
@@ -216,9 +215,7 @@ func (o *buildOptions) toDisplayMode() (progressui.DisplayMode, error) {
 	return progress, nil
 }
 
-func runBuild(dockerCli command.Cli, options buildOptions) (err error) {
-	ctx := appcontext.Context()
-
+func runBuild(ctx context.Context, dockerCli command.Cli, options buildOptions) (err error) {
 	mp, report, err := metrics.MeterProvider(dockerCli)
 	if err != nil {
 		return err
@@ -487,7 +484,7 @@ func buildCmd(dockerCli command.Cli, rootOpts *rootOptions, debugConfig *debug.D
 				options.invokeConfig = iConfig
 			}
 
-			return runBuild(dockerCli, *options)
+			return runBuild(cmd.Context(), dockerCli, *options)
 		},
 		ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 			return nil, cobra.ShellCompDirectiveFilterDirs
diff --git a/commands/create.go b/commands/create.go
index f9f847c2b96b..9362aebba643 100644
--- a/commands/create.go
+++ b/commands/create.go
@@ -2,6 +2,7 @@ package commands
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	"github.com/docker/buildx/builder"
@@ -11,7 +12,6 @@ import (
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/spf13/cobra"
 )
 
@@ -30,9 +30,7 @@ type createOptions struct {
 	// upgrade      bool // perform upgrade of the driver
 }
 
-func runCreate(dockerCli command.Cli, in createOptions, args []string) error {
-	ctx := appcontext.Context()
-
+func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, args []string) error {
 	txn, release, err := storeutil.GetStore(dockerCli)
 	if err != nil {
 		return err
@@ -98,7 +96,7 @@ func createCmd(dockerCli command.Cli) *cobra.Command {
 		Short: "Create a new builder instance",
 		Args:  cli.RequiresMaxArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runCreate(dockerCli, options, args)
+			return runCreate(cmd.Context(), dockerCli, options, args)
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/diskusage.go b/commands/diskusage.go
index dd25e3f3c0e8..4460c0fee38f 100644
--- a/commands/diskusage.go
+++ b/commands/diskusage.go
@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -15,7 +16,6 @@ import (
 	"github.com/docker/cli/opts"
 	"github.com/docker/go-units"
 	"github.com/moby/buildkit/client"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
 )
@@ -26,9 +26,7 @@ type duOptions struct {
 	verbose bool
 }
 
-func runDiskUsage(dockerCli command.Cli, opts duOptions) error {
-	ctx := appcontext.Context()
-
+func runDiskUsage(ctx context.Context, dockerCli command.Cli, opts duOptions) error {
 	pi, err := toBuildkitPruneInfo(opts.filter.Value())
 	if err != nil {
 		return err
@@ -114,7 +112,7 @@ func duCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Args:  cli.NoArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
-			return runDiskUsage(dockerCli, options)
+			return runDiskUsage(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/imagetools/create.go b/commands/imagetools/create.go
index 0222411312e6..f05571b03942 100644
--- a/commands/imagetools/create.go
+++ b/commands/imagetools/create.go
@@ -13,7 +13,6 @@ import (
 	"github.com/docker/buildx/util/imagetools"
 	"github.com/docker/buildx/util/progress"
 	"github.com/docker/cli/cli/command"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/moby/buildkit/util/progress/progressui"
 	"github.com/opencontainers/go-digest"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
@@ -32,7 +31,7 @@ type createOptions struct {
 	progress     string
 }
 
-func runCreate(dockerCli command.Cli, in createOptions, args []string) error {
+func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, args []string) error {
 	if len(args) == 0 && len(in.files) == 0 {
 		return errors.Errorf("no sources specified")
 	}
@@ -113,8 +112,6 @@ func runCreate(dockerCli command.Cli, in createOptions, args []string) error {
 		}
 	}
 
-	ctx := appcontext.Context()
-
 	b, err := builder.New(dockerCli, builder.WithName(in.builder))
 	if err != nil {
 		return err
@@ -274,7 +271,7 @@ func createCmd(dockerCli command.Cli, opts RootOptions) *cobra.Command {
 		Short: "Create a new image based on source images",
 		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = *opts.Builder
-			return runCreate(dockerCli, options, args)
+			return runCreate(cmd.Context(), dockerCli, options, args)
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/imagetools/inspect.go b/commands/imagetools/inspect.go
index ad1e6438a662..5b06f52fc63c 100644
--- a/commands/imagetools/inspect.go
+++ b/commands/imagetools/inspect.go
@@ -1,13 +1,14 @@
 package commands
 
 import (
+	"context"
+
 	"github.com/docker/buildx/builder"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/buildx/util/imagetools"
 	"github.com/docker/cli-docs-tool/annotation"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 )
@@ -18,9 +19,7 @@ type inspectOptions struct {
 	raw     bool
 }
 
-func runInspect(dockerCli command.Cli, in inspectOptions, name string) error {
-	ctx := appcontext.Context()
-
+func runInspect(ctx context.Context, dockerCli command.Cli, in inspectOptions, name string) error {
 	if in.format != "" && in.raw {
 		return errors.Errorf("format and raw cannot be used together")
 	}
@@ -51,7 +50,7 @@ func inspectCmd(dockerCli command.Cli, rootOpts RootOptions) *cobra.Command {
 		Args:  cli.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = *rootOpts.Builder
-			return runInspect(dockerCli, options, args[0])
+			return runInspect(cmd.Context(), dockerCli, options, args[0])
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/inspect.go b/commands/inspect.go
index bae2ac6d6580..e341691c64d8 100644
--- a/commands/inspect.go
+++ b/commands/inspect.go
@@ -17,7 +17,6 @@ import (
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/debug"
 	"github.com/docker/go-units"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/spf13/cobra"
 )
 
@@ -26,9 +25,7 @@ type inspectOptions struct {
 	builder   string
 }
 
-func runInspect(dockerCli command.Cli, in inspectOptions) error {
-	ctx := appcontext.Context()
-
+func runInspect(ctx context.Context, dockerCli command.Cli, in inspectOptions) error {
 	b, err := builder.New(dockerCli,
 		builder.WithName(in.builder),
 		builder.WithSkippedValidation(),
@@ -150,7 +147,7 @@ func inspectCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 			if len(args) > 0 {
 				options.builder = args[0]
 			}
-			return runInspect(dockerCli, options)
+			return runInspect(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
diff --git a/commands/ls.go b/commands/ls.go
index 0be95f7ab04c..3b18e42cc779 100644
--- a/commands/ls.go
+++ b/commands/ls.go
@@ -17,7 +17,6 @@ import (
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/command/formatter"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
 )
@@ -39,9 +38,7 @@ type lsOptions struct {
 	format string
 }
 
-func runLs(dockerCli command.Cli, in lsOptions) error {
-	ctx := appcontext.Context()
-
+func runLs(ctx context.Context, dockerCli command.Cli, in lsOptions) error {
 	txn, release, err := storeutil.GetStore(dockerCli)
 	if err != nil {
 		return err
@@ -103,7 +100,7 @@ func lsCmd(dockerCli command.Cli) *cobra.Command {
 		Short: "List builder instances",
 		Args:  cli.ExactArgs(0),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			return runLs(dockerCli, options)
+			return runLs(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/prune.go b/commands/prune.go
index 454b04a8c5e9..1a2a5220c02d 100644
--- a/commands/prune.go
+++ b/commands/prune.go
@@ -1,6 +1,7 @@
 package commands
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -15,7 +16,6 @@ import (
 	"github.com/docker/docker/api/types/filters"
 	"github.com/docker/go-units"
 	"github.com/moby/buildkit/client"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
@@ -35,9 +35,7 @@ const (
 	allCacheWarning = `WARNING! This will remove all build cache. Are you sure you want to continue?`
 )
 
-func runPrune(dockerCli command.Cli, opts pruneOptions) error {
-	ctx := appcontext.Context()
-
+func runPrune(ctx context.Context, dockerCli command.Cli, opts pruneOptions) error {
 	pruneFilters := opts.filter.Value()
 	pruneFilters = command.PruneFilters(dockerCli, pruneFilters)
 
@@ -138,7 +136,7 @@ func pruneCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 		Args:  cli.NoArgs,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			options.builder = rootOpts.builder
-			return runPrune(dockerCli, options)
+			return runPrune(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.Disable,
 	}
diff --git a/commands/rm.go b/commands/rm.go
index 65293ec69951..af6ffd59f190 100644
--- a/commands/rm.go
+++ b/commands/rm.go
@@ -10,7 +10,6 @@ import (
 	"github.com/docker/buildx/store/storeutil"
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli/command"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/pkg/errors"
 	"github.com/spf13/cobra"
 	"golang.org/x/sync/errgroup"
@@ -28,9 +27,7 @@ const (
 	rmInactiveWarning = `WARNING! This will remove all builders that are not in running state. Are you sure you want to continue?`
 )
 
-func runRm(dockerCli command.Cli, in rmOptions) error {
-	ctx := appcontext.Context()
-
+func runRm(ctx context.Context, dockerCli command.Cli, in rmOptions) error {
 	if in.allInactive && !in.force && !command.PromptForConfirmation(dockerCli.In(), dockerCli.Out(), rmInactiveWarning) {
 		return nil
 	}
@@ -108,7 +105,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 				}
 				options.builders = args
 			}
-			return runRm(dockerCli, options)
+			return runRm(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}
diff --git a/commands/root.go b/commands/root.go
index 780b689d1d72..1c862489331b 100644
--- a/commands/root.go
+++ b/commands/root.go
@@ -13,6 +13,7 @@ import (
 	"github.com/docker/cli/cli-plugins/plugin"
 	"github.com/docker/cli/cli/command"
 	"github.com/docker/cli/cli/debug"
+	"github.com/moby/buildkit/util/appcontext"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -29,12 +30,15 @@ func NewRootCmd(name string, isPlugin bool, dockerCli command.Cli) *cobra.Comman
 		CompletionOptions: cobra.CompletionOptions{
 			HiddenDefaultCmd: true,
 		},
-	}
-	if isPlugin {
-		cmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error {
+		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
+			cmd.SetContext(appcontext.Context())
+			if !isPlugin {
+				return nil
+			}
 			return plugin.PersistentPreRunE(cmd, args)
-		}
-	} else {
+		},
+	}
+	if !isPlugin {
 		// match plugin behavior for standalone mode
 		// https://github.com/docker/cli/blob/6c9eb708fa6d17765d71965f90e1c59cea686ee9/cli-plugins/plugin/plugin.go#L117-L127
 		cmd.SilenceUsage = true
diff --git a/commands/stop.go b/commands/stop.go
index c91b8ad39202..b27985af8a24 100644
--- a/commands/stop.go
+++ b/commands/stop.go
@@ -7,7 +7,6 @@ import (
 	"github.com/docker/buildx/util/cobrautil/completion"
 	"github.com/docker/cli/cli"
 	"github.com/docker/cli/cli/command"
-	"github.com/moby/buildkit/util/appcontext"
 	"github.com/spf13/cobra"
 )
 
@@ -15,9 +14,7 @@ type stopOptions struct {
 	builder string
 }
 
-func runStop(dockerCli command.Cli, in stopOptions) error {
-	ctx := appcontext.Context()
-
+func runStop(ctx context.Context, dockerCli command.Cli, in stopOptions) error {
 	b, err := builder.New(dockerCli,
 		builder.WithName(in.builder),
 		builder.WithSkippedValidation(),
@@ -45,7 +42,7 @@ func stopCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 			if len(args) > 0 {
 				options.builder = args[0]
 			}
-			return runStop(dockerCli, options)
+			return runStop(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
 	}