cmd/cue: give access to cue Command in run functions

This allows more control over the overall error state.

Issue #101
Issue #100

Change-Id: Icaf81bd8757bfed6522ec0aff249afa7706d8495
Reviewed-on: https://cue-review.googlesource.com/c/cue/+/3101
Reviewed-by: Marcel van Lohuizen <mpvl@google.com>
diff --git a/cmd/cue/cmd/add.go b/cmd/cue/cmd/add.go
index 559b911..a4b2b3e 100644
--- a/cmd/cue/cmd/add.go
+++ b/cmd/cue/cmd/add.go
@@ -34,7 +34,7 @@
 	"github.com/spf13/cobra"
 )
 
-func newAddCmd() *cobra.Command {
+func newAddCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		// TODO: this command is still experimental, don't show it in
 		// the documentation just yet.
@@ -44,7 +44,7 @@
 		Short: "bulk append to CUE files",
 		Long: `Append a common snippet of CUE to many files and commit atomically.
 `,
-		RunE: runAdd,
+		RunE: mkRunE(c, runAdd),
 	}
 
 	f := cmd.Flags()
@@ -56,13 +56,13 @@
 	return cmd
 }
 
-func runAdd(cmd *cobra.Command, args []string) (err error) {
+func runAdd(cmd *Command, args []string) (err error) {
 	return doAdd(cmd, stdin, args)
 }
 
 var stdin io.Reader = os.Stdin
 
-func doAdd(cmd *cobra.Command, stdin io.Reader, args []string) (err error) {
+func doAdd(cmd *Command, stdin io.Reader, args []string) (err error) {
 	// Offsets at which to restore original files, if any, if any of the
 	// appends fail.
 	// Ideally this is placed below where it is used, but we want to make
@@ -192,7 +192,7 @@
 	contents []byte
 }
 
-func restoreOriginals(cmd *cobra.Command, originals []originalFile) {
+func restoreOriginals(cmd *Command, originals []originalFile) {
 	for _, fo := range originals {
 		if err := fo.restore(); err != nil {
 			fmt.Fprintln(cmd.OutOrStderr(), "Error restoring file: ", err)
@@ -214,7 +214,7 @@
 	build    *build.Instance
 }
 
-func initFile(cmd *cobra.Command, file string, getBuild func(path string) *build.Instance) (todo *fileInfo, err error) {
+func initFile(cmd *Command, file string, getBuild func(path string) *build.Instance) (todo *fileInfo, err error) {
 	defer func() {
 		if err != nil {
 			err = fmt.Errorf("init file: %v", err)
diff --git a/cmd/cue/cmd/cmd.go b/cmd/cue/cmd/cmd.go
index ad466c5..9e5fb0c 100644
--- a/cmd/cue/cmd/cmd.go
+++ b/cmd/cue/cmd/cmd.go
@@ -23,7 +23,7 @@
 
 // TODO: generate long description from documentation.
 
-func newCmdCmd() *cobra.Command {
+func newCmdCmd(c *Command) *cobra.Command {
 	return &cobra.Command{
 		Use:   "cmd <name> [-x] [instances]",
 		Short: "run a user-defined shell command",
@@ -210,7 +210,7 @@
 	}
 
 `,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: mkRunE(c, func(cmd *Command, args []string) error {
 			if len(args) == 0 {
 				fmt.Println("cmd must be run as one of its subcommands")
 			} else {
@@ -219,6 +219,6 @@
 			fmt.Println("Run 'cue help cmd' for known subcommands.")
 			os.Exit(1) // TODO: get rid of this
 			return nil
-		},
+		}),
 	}
 }
diff --git a/cmd/cue/cmd/cmd_test.go b/cmd/cue/cmd/cmd_test.go
index c24e638..d09c41f 100644
--- a/cmd/cue/cmd/cmd_test.go
+++ b/cmd/cue/cmd/cmd_test.go
@@ -38,13 +38,13 @@
 		stderr = os.Stderr
 	}()
 	for _, name := range testCases {
-		rootCmd := newRootCmd().root
+		c := newRootCmd()
 		run := func(cmd *cobra.Command, args []string) error {
 			stdout = cmd.OutOrStdout()
 			stderr = cmd.OutOrStderr()
 
-			tools, _ := buildTools(rootCmd, args)
-			cmd, err := addCustom(rootCmd, "command", name, tools)
+			tools, _ := buildTools(c, args)
+			cmd, err := addCustom(c, c.root, "command", name, tools)
 			if err != nil {
 				return err
 			}
diff --git a/cmd/cue/cmd/common.go b/cmd/cue/cmd/common.go
index 9908328..2bf9f37 100644
--- a/cmd/cue/cmd/common.go
+++ b/cmd/cue/cmd/common.go
@@ -40,7 +40,7 @@
 	}
 }
 
-func exitIfErr(cmd *cobra.Command, inst *cue.Instance, err error, fatal bool) {
+func exitIfErr(cmd *Command, inst *cue.Instance, err error, fatal bool) {
 	exitOnErr(cmd, err, fatal)
 }
 
@@ -53,7 +53,7 @@
 	return language.Make(loc)
 }
 
-func exitOnErr(cmd *cobra.Command, err error, fatal bool) {
+func exitOnErr(cmd *Command, err error, fatal bool) {
 	if err == nil {
 		return
 	}
@@ -80,7 +80,7 @@
 	}
 }
 
-func buildFromArgs(cmd *cobra.Command, args []string) []*cue.Instance {
+func buildFromArgs(cmd *Command, args []string) []*cue.Instance {
 	binst := loadFromArgs(cmd, args)
 	if binst == nil {
 		return nil
@@ -88,7 +88,7 @@
 	return buildInstances(cmd, binst)
 }
 
-func loadFromArgs(cmd *cobra.Command, args []string) []*build.Instance {
+func loadFromArgs(cmd *Command, args []string) []*build.Instance {
 	log.SetOutput(cmd.OutOrStderr())
 	binst := load.Instances(args, nil)
 	if len(binst) == 0 {
@@ -97,7 +97,7 @@
 	return binst
 }
 
-func buildInstances(cmd *cobra.Command, binst []*build.Instance) []*cue.Instance {
+func buildInstances(cmd *Command, binst []*build.Instance) []*cue.Instance {
 	instances := cue.Build(binst)
 	for _, inst := range instances {
 		// TODO: consider merging errors of multiple files, but ensure
@@ -118,7 +118,7 @@
 	return instances
 }
 
-func buildToolInstances(cmd *cobra.Command, binst []*build.Instance) ([]*cue.Instance, error) {
+func buildToolInstances(cmd *Command, binst []*build.Instance) ([]*cue.Instance, error) {
 	instances := cue.Build(binst)
 	for _, inst := range instances {
 		if inst.Err != nil {
@@ -135,7 +135,7 @@
 	return instances, nil
 }
 
-func buildTools(cmd *cobra.Command, args []string) (*cue.Instance, error) {
+func buildTools(cmd *Command, args []string) (*cue.Instance, error) {
 	binst := loadFromArgs(cmd, args)
 	if len(binst) == 0 {
 		return nil, nil
diff --git a/cmd/cue/cmd/common_test.go b/cmd/cue/cmd/common_test.go
index 9d8c73b..58333d9 100644
--- a/cmd/cue/cmd/common_test.go
+++ b/cmd/cue/cmd/common_test.go
@@ -121,5 +121,5 @@
 }
 
 func TestLoadError(t *testing.T) {
-	runCommand(t, newEvalCmd(), "loaderr", "non-existing", ".")
+	runCommand(t, newEvalCmd(newRootCmd()), "loaderr", "non-existing", ".")
 }
diff --git a/cmd/cue/cmd/custom.go b/cmd/cue/cmd/custom.go
index 63eb3d8..f04dc94 100644
--- a/cmd/cue/cmd/custom.go
+++ b/cmd/cue/cmd/custom.go
@@ -58,7 +58,7 @@
 	stderr io.Writer = os.Stderr
 )
 
-func addCustom(parent *cobra.Command, typ, name string, tools *cue.Instance) (*cobra.Command, error) {
+func addCustom(c *Command, parent *cobra.Command, typ, name string, tools *cue.Instance) (*cobra.Command, error) {
 	if tools == nil {
 		return nil, errors.New("no commands defined")
 	}
@@ -77,13 +77,13 @@
 		Use:   usage,
 		Short: lookupString(o, "short"),
 		Long:  lookupString(o, "long"),
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: mkRunE(c, func(cmd *Command, args []string) error {
 			// TODO:
 			// - parse flags and env vars
 			// - constrain current config with config section
 
 			return doTasks(cmd, typ, name, tools)
-		},
+		}),
 	}
 	parent.AddCommand(sub)
 
@@ -121,7 +121,7 @@
 	return root.Lookup(k.typ, k.name, taskSection)
 }
 
-func doTasks(cmd *cobra.Command, typ, command string, root *cue.Instance) error {
+func doTasks(cmd *Command, typ, command string, root *cue.Instance) error {
 	err := executeTasks(typ, command, root)
 	exitIfErr(cmd, root, err, true)
 	return err
diff --git a/cmd/cue/cmd/eval.go b/cmd/cue/cmd/eval.go
index a8e667d..bd70b73 100644
--- a/cmd/cue/cmd/eval.go
+++ b/cmd/cue/cmd/eval.go
@@ -25,7 +25,7 @@
 )
 
 // newEvalCmd creates a new eval command
-func newEvalCmd() *cobra.Command {
+func newEvalCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "eval",
 		Short: "evaluate and print a configuration",
@@ -46,7 +46,7 @@
   "a"
   "c"
 `,
-		RunE: runEval,
+		RunE: mkRunE(c, runEval),
 	}
 
 	cmd.Flags().StringArrayP(string(flagExpression), "e", nil, "evaluate this expression only")
@@ -77,7 +77,7 @@
 	flagAttributes flagName = "attributes"
 )
 
-func runEval(cmd *cobra.Command, args []string) error {
+func runEval(cmd *Command, args []string) error {
 	instances := buildFromArgs(cmd, args)
 
 	var exprs []ast.Expr
diff --git a/cmd/cue/cmd/eval_test.go b/cmd/cue/cmd/eval_test.go
index 1d7d20a..56c4c5a 100644
--- a/cmd/cue/cmd/eval_test.go
+++ b/cmd/cue/cmd/eval_test.go
@@ -17,13 +17,13 @@
 import "testing"
 
 func TestEval(t *testing.T) {
-	runCommand(t, newEvalCmd(), "eval")
+	runCommand(t, newEvalCmd(newRootCmd()), "eval")
 
-	cmd := newEvalCmd()
+	cmd := newEvalCmd(newRootCmd())
 	mustParseFlags(t, cmd, "-c", "-a")
 	runCommand(t, cmd, "eval_conc")
 
-	cmd = newEvalCmd()
+	cmd = newEvalCmd(newRootCmd())
 	mustParseFlags(t, cmd, "-c", "-e", "b.a.b", "-e", "b.idx")
 	runCommand(t, cmd, "eval_expr")
 }
diff --git a/cmd/cue/cmd/export.go b/cmd/cue/cmd/export.go
index beee972..783c983 100644
--- a/cmd/cue/cmd/export.go
+++ b/cmd/cue/cmd/export.go
@@ -24,7 +24,7 @@
 )
 
 // newExportCmd creates and export command
-func newExportCmd() *cobra.Command {
+func newExportCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "export",
 		Short: "output data in a standard format",
@@ -86,7 +86,7 @@
         The evaluated value must be of type string.
 `,
 
-		RunE: runExport,
+		RunE: mkRunE(c, runExport),
 	}
 	flagMedia.Add(cmd)
 	cmd.Flags().Bool(string(flagEscape), false, "use HTML escaping")
@@ -94,7 +94,7 @@
 	return cmd
 }
 
-func runExport(cmd *cobra.Command, args []string) error {
+func runExport(cmd *Command, args []string) error {
 	instances := buildFromArgs(cmd, args)
 	w := cmd.OutOrStdout()
 
@@ -117,7 +117,7 @@
 	return nil
 }
 
-func outputJSON(cmd *cobra.Command, w io.Writer, v cue.Value) error {
+func outputJSON(cmd *Command, w io.Writer, v cue.Value) error {
 	e := json.NewEncoder(w)
 	e.SetIndent("", "    ")
 	e.SetEscapeHTML(flagEscape.Bool(cmd))
diff --git a/cmd/cue/cmd/export_test.go b/cmd/cue/cmd/export_test.go
index 9067630..9b9aac5 100644
--- a/cmd/cue/cmd/export_test.go
+++ b/cmd/cue/cmd/export_test.go
@@ -17,6 +17,6 @@
 import "testing"
 
 func TestExport(t *testing.T) {
-	runCommand(t, newExportCmd(), "export")
-	runCommand(t, newExportCmd(), "export_err")
+	runCommand(t, newExportCmd(newRootCmd()), "export")
+	runCommand(t, newExportCmd(newRootCmd()), "export_err")
 }
diff --git a/cmd/cue/cmd/flags.go b/cmd/cue/cmd/flags.go
index 815ec79..f202eab 100644
--- a/cmd/cue/cmd/flags.go
+++ b/cmd/cue/cmd/flags.go
@@ -69,17 +69,17 @@
 
 type flagName string
 
-func (f flagName) Bool(cmd *cobra.Command) bool {
+func (f flagName) Bool(cmd *Command) bool {
 	v, _ := cmd.Flags().GetBool(string(f))
 	return v
 }
 
-func (f flagName) String(cmd *cobra.Command) string {
+func (f flagName) String(cmd *Command) string {
 	v, _ := cmd.Flags().GetString(string(f))
 	return v
 }
 
-func (f flagName) StringArray(cmd *cobra.Command) []string {
+func (f flagName) StringArray(cmd *Command) []string {
 	v, _ := cmd.Flags().GetStringArray(string(f))
 	return v
 }
@@ -95,7 +95,7 @@
 	cmd.Flags().StringP(f.name, f.short, f.def, f.text)
 }
 
-func (f *stringFlag) String(cmd *cobra.Command) string {
+func (f *stringFlag) String(cmd *Command) string {
 	v, err := cmd.Flags().GetString(f.name)
 	if err != nil {
 		return f.def
diff --git a/cmd/cue/cmd/fmt.go b/cmd/cue/cmd/fmt.go
index b2f1b75..93298a2 100644
--- a/cmd/cue/cmd/fmt.go
+++ b/cmd/cue/cmd/fmt.go
@@ -23,13 +23,13 @@
 	"github.com/spf13/cobra"
 )
 
-func newFmtCmd() *cobra.Command {
+func newFmtCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "fmt [-s] [packages]",
 		Short: "formats CUE configuration files",
 		Long: `Fmt formats the given files or the files for the given packages in place
 `,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: mkRunE(c, func(cmd *Command, args []string) error {
 			for _, inst := range load.Instances(args, nil) {
 				all := []string{}
 				all = append(all, inst.CUEFiles...)
@@ -65,7 +65,7 @@
 				}
 			}
 			return nil
-		},
+		}),
 	}
 	return cmd
 }
diff --git a/cmd/cue/cmd/get.go b/cmd/cue/cmd/get.go
index 6cf6049..499e62c 100644
--- a/cmd/cue/cmd/get.go
+++ b/cmd/cue/cmd/get.go
@@ -21,7 +21,7 @@
 	"github.com/spf13/cobra"
 )
 
-func newGetCmd() *cobra.Command {
+func newGetCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "get <language> [packages]",
 		Short: "add dependencies to the current module",
@@ -35,7 +35,7 @@
 The specifics on how dependencies are fechted and converted vary
 per language and are documented in the respective subcommands.
 `,
-		RunE: func(cmd *cobra.Command, args []string) error {
+		RunE: mkRunE(c, func(cmd *Command, args []string) error {
 			if len(args) == 0 {
 				fmt.Println("get must be run as one of its subcommands")
 			} else {
@@ -44,8 +44,8 @@
 			fmt.Println("Run 'cue help get' for known subcommands.")
 			os.Exit(1) // TODO: get rid of this
 			return nil
-		},
+		}),
 	}
-	cmd.AddCommand(newGoCmd())
+	cmd.AddCommand(newGoCmd(c))
 	return cmd
 }
diff --git a/cmd/cue/cmd/get_go.go b/cmd/cue/cmd/get_go.go
index f9f5821..411cf71 100644
--- a/cmd/cue/cmd/get_go.go
+++ b/cmd/cue/cmd/get_go.go
@@ -50,7 +50,7 @@
 //   package foo
 //   Type: enumType
 
-func newGoCmd() *cobra.Command {
+func newGoCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "go [packages]",
 		Short: "add Go dependencies to the current module",
@@ -201,9 +201,7 @@
 `,
 		// - TODO: interpret cuego's struct tags and annotations.
 
-		RunE: func(cmd *cobra.Command, args []string) error {
-			return extract(cmd, args)
-		},
+		RunE: mkRunE(c, extract),
 	}
 
 	cmd.Flags().StringP(string(flagExclude), "e", "",
@@ -240,7 +238,7 @@
 }
 
 type extractor struct {
-	cmd *cobra.Command
+	cmd *Command
 
 	stderr io.Writer
 	pkgs   []*packages.Package
@@ -313,7 +311,7 @@
 // TODO:
 // - consider not including types with any dropped fields.
 
-func extract(cmd *cobra.Command, args []string) error {
+func extract(cmd *Command, args []string) error {
 	// determine module root:
 	binst := loadFromArgs(cmd, []string{"."})[0]
 
diff --git a/cmd/cue/cmd/get_go_test.go b/cmd/cue/cmd/get_go_test.go
index 577ab8a..509be39 100644
--- a/cmd/cue/cmd/get_go_test.go
+++ b/cmd/cue/cmd/get_go_test.go
@@ -38,7 +38,7 @@
 	cueTestRoot = tmp
 
 	// We don't use runCommand here, as we are interested in generated packages.
-	cmd := newGoCmd()
+	cmd := newGoCmd(newRootCmd())
 	cmd.SetArgs([]string{"./testdata/code/go/..."})
 	err = cmd.Execute()
 	if err != nil {
diff --git a/cmd/cue/cmd/import.go b/cmd/cue/cmd/import.go
index 52df22e..612f64d 100644
--- a/cmd/cue/cmd/import.go
+++ b/cmd/cue/cmd/import.go
@@ -42,7 +42,7 @@
 	"golang.org/x/sync/errgroup"
 )
 
-func newImportCmd() *cobra.Command {
+func newImportCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "import",
 		Short: "convert other data formats to CUE files",
@@ -193,7 +193,7 @@
       }
   }
 `,
-		RunE: runImport,
+		RunE: mkRunE(c, runImport),
 	}
 
 	flagOut.Add(cmd)
@@ -220,7 +220,7 @@
 )
 
 type importStreamFunc func(path string, r io.Reader) ([]ast.Expr, error)
-type importFileFunc func(cmd *cobra.Command, path string, r io.Reader) (*ast.File, error)
+type importFileFunc func(cmd *Command, path string, r io.Reader) (*ast.File, error)
 
 type encodingInfo struct {
 	fnStream importStreamFunc
@@ -250,7 +250,7 @@
 	return nil
 }
 
-func runImport(cmd *cobra.Command, args []string) error {
+func runImport(cmd *Command, args []string) error {
 	log.SetOutput(cmd.OutOrStderr())
 
 	var group errgroup.Group
@@ -307,7 +307,7 @@
 	return nil
 }
 
-func handleFile(cmd *cobra.Command, pkg, filename string) (err error) {
+func handleFile(cmd *Command, pkg, filename string) (err error) {
 	re, err := regexp.Compile(flagGlob.String(cmd))
 	if err != nil {
 		return err
@@ -348,7 +348,7 @@
 	}
 }
 
-func processFile(cmd *cobra.Command, file *ast.File) (err error) {
+func processFile(cmd *Command, file *ast.File) (err error) {
 	name := file.Filename + ".cue"
 
 	b, err := format.Node(file)
@@ -359,7 +359,7 @@
 	return ioutil.WriteFile(name, b, 0644)
 }
 
-func processStream(cmd *cobra.Command, pkg, filename string, objs []ast.Expr) error {
+func processStream(cmd *Command, pkg, filename string, objs []ast.Expr) error {
 	if flagFiles.Bool(cmd) {
 		for i, f := range objs {
 			err := combineExpressions(cmd, pkg, newName(filename, i), f)
@@ -379,7 +379,7 @@
 // TODO: implement a more fine-grained approach.
 var mutex sync.Mutex
 
-func combineExpressions(cmd *cobra.Command, pkg, cueFile string, objs ...ast.Expr) error {
+func combineExpressions(cmd *Command, pkg, cueFile string, objs ...ast.Expr) error {
 	mutex.Lock()
 	defer mutex.Unlock()
 
@@ -640,7 +640,7 @@
 	return objects, nil
 }
 
-func handleProtoDef(cmd *cobra.Command, path string, r io.Reader) (f *ast.File, err error) {
+func handleProtoDef(cmd *Command, path string, r io.Reader) (f *ast.File, err error) {
 	return protobuf.Extract(path, r, &protobuf.Config{Paths: flagProtoPath.StringArray(cmd)})
 }
 
diff --git a/cmd/cue/cmd/import_test.go b/cmd/cue/cmd/import_test.go
index 9d8d3cd..c892fbd 100644
--- a/cmd/cue/cmd/import_test.go
+++ b/cmd/cue/cmd/import_test.go
@@ -19,22 +19,22 @@
 )
 
 func TestImport(t *testing.T) {
-	cmd := newImportCmd()
+	cmd := newImportCmd(newRootCmd())
 	mustParseFlags(t, cmd,
 		"-o", "-", "-f", "--files")
 	runCommand(t, cmd, "import_files")
 
-	cmd = newImportCmd()
+	cmd = newImportCmd(newRootCmd())
 	mustParseFlags(t, cmd,
 		"-o", "-", "-f", "-l", `"\(strings.ToLower(kind))" "\(name)"`)
 	runCommand(t, cmd, "import_path")
 
-	cmd = newImportCmd()
+	cmd = newImportCmd(newRootCmd())
 	mustParseFlags(t, cmd,
 		"-o", "-", "-f", "-l", `"\(strings.ToLower(kind))"`, "--list")
 	runCommand(t, cmd, "import_list")
 
-	cmd = newImportCmd()
+	cmd = newImportCmd(newRootCmd())
 	mustParseFlags(t, cmd,
 		"-o", "-", "-f", "--list",
 		"-l", `"\(strings.ToLower(kind))" "\(name)"`, "--recursive")
diff --git a/cmd/cue/cmd/root.go b/cmd/cue/cmd/root.go
index e34a40a..433ffe4 100644
--- a/cmd/cue/cmd/root.go
+++ b/cmd/cue/cmd/root.go
@@ -41,6 +41,15 @@
 
 var log = logger.New(os.Stderr, "", logger.Lshortfile)
 
+type runFunction func(cmd *Command, args []string) error
+
+func mkRunE(c *Command, f runFunction) func(*cobra.Command, []string) error {
+	return func(cmd *cobra.Command, args []string) error {
+		c.Command = cmd
+		return f(c, args)
+	}
+}
+
 // newRootCmd creates the base command when called without any subcommands
 func newRootCmd() *Command {
 	cmd := &cobra.Command{
@@ -69,19 +78,22 @@
 		SilenceUsage: true,
 	}
 
-	cmdCmd := newCmdCmd()
+	c := &Command{Command: cmd, root: cmd}
+
+	cmdCmd := newCmdCmd(c)
+	c.cmd = cmdCmd
 
 	subCommands := []*cobra.Command{
-		newTrimCmd(),
-		newImportCmd(),
-		newEvalCmd(),
-		newGetCmd(),
-		newFmtCmd(),
-		newExportCmd(),
+		newTrimCmd(c),
+		newImportCmd(c),
+		newEvalCmd(c),
+		newGetCmd(c),
+		newFmtCmd(c),
+		newExportCmd(c),
 		cmdCmd,
-		newVersionCmd(),
-		newVetCmd(),
-		newAddCmd(),
+		newVersionCmd(c),
+		newVetCmd(c),
+		newAddCmd(c),
 	}
 
 	addGlobalFlags(cmd.PersistentFlags())
@@ -90,7 +102,7 @@
 		cmd.AddCommand(sub)
 	}
 
-	return &Command{root: cmd, cmd: cmdCmd}
+	return c
 }
 
 // Main runs the cue tool. It loads the tool flags.
@@ -106,6 +118,9 @@
 }
 
 type Command struct {
+	// The currently active command.
+	*cobra.Command
+
 	root *cobra.Command
 
 	// Subcommands
@@ -178,11 +193,11 @@
 		return cmd, nil // Forces unknown command message from Cobra.
 	}
 
-	tools, err := buildTools(rootCmd, args[1:])
+	tools, err := buildTools(cmd, args[1:])
 	if err != nil {
 		return cmd, err
 	}
-	_, err = addCustom(rootCmd, commandSection, args[0], tools)
+	_, err = addCustom(cmd, rootCmd, commandSection, args[0], tools)
 	if err != nil {
 		fmt.Printf("command %s %q is not defined\n", commandSection, args[0])
 		fmt.Println("Run 'cue help' to show available commands.")
@@ -212,7 +227,7 @@
 		args = args[1:]
 	}
 
-	tools, err := buildTools(cmd.root, args)
+	tools, err := buildTools(cmd, args)
 	if err != nil {
 		return err
 	}
@@ -230,7 +245,7 @@
 			return errors.Newf(token.NoPos, "could not create command definitions: %v", err)
 		}
 		for i.Next() {
-			_, _ = addCustom(spec.cmd, spec.name, i.Label(), tools)
+			_, _ = addCustom(cmd, spec.cmd, spec.name, i.Label(), tools)
 		}
 	}
 	return nil
diff --git a/cmd/cue/cmd/trim.go b/cmd/cue/cmd/trim.go
index 1ca999f..0e13c26 100644
--- a/cmd/cue/cmd/trim.go
+++ b/cmd/cue/cmd/trim.go
@@ -35,7 +35,7 @@
 // - implement verification post-processing as extra safety
 
 // newTrimCmd creates a trim command
-func newTrimCmd() *cobra.Command {
+func newTrimCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "trim",
 		Short: "remove superfluous fields",
@@ -84,14 +84,14 @@
 It is guaranteed that the resulting files give the same output as before the
 removal.
 `,
-		RunE: runTrim,
+		RunE: mkRunE(c, runTrim),
 	}
 
 	flagOut.Add(cmd)
 	return cmd
 }
 
-func runTrim(cmd *cobra.Command, args []string) error {
+func runTrim(cmd *Command, args []string) error {
 	// TODO: Do something more fine-grained. Optional fields are mostly not
 	// useful to consider as an optional field will almost never subsume
 	// another value. However, an optional field may subsume and therefore
@@ -167,14 +167,14 @@
 }
 
 type trimSet struct {
-	cmd       *cobra.Command
+	cmd       *Command
 	stack     []string
 	exclude   map[ast.Node]bool // don't remove fields marked here
 	alwaysGen map[ast.Node]bool // node is always from a generated source
 	fromComp  map[ast.Node]bool // node originated from a comprehension
 }
 
-func newTrimSet(cmd *cobra.Command) *trimSet {
+func newTrimSet(cmd *Command) *trimSet {
 	return &trimSet{
 		cmd:       cmd,
 		exclude:   map[ast.Node]bool{},
diff --git a/cmd/cue/cmd/trim_test.go b/cmd/cue/cmd/trim_test.go
index bc8001f..6326f31 100644
--- a/cmd/cue/cmd/trim_test.go
+++ b/cmd/cue/cmd/trim_test.go
@@ -19,7 +19,7 @@
 )
 
 func TestTrim(t *testing.T) {
-	cmd := newTrimCmd()
+	cmd := newTrimCmd(newRootCmd())
 	mustParseFlags(t, cmd, "-o", "-")
 	runCommand(t, cmd, "trim")
 }
diff --git a/cmd/cue/cmd/version.go b/cmd/cue/cmd/version.go
index d605469..be58c82 100644
--- a/cmd/cue/cmd/version.go
+++ b/cmd/cue/cmd/version.go
@@ -21,12 +21,12 @@
 	"github.com/spf13/cobra"
 )
 
-func newVersionCmd() *cobra.Command {
+func newVersionCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "version",
 		Short: "print CUE version",
 		Long:  ``,
-		RunE:  runVersion,
+		RunE:  mkRunE(c, runVersion),
 	}
 	return cmd
 }
@@ -37,7 +37,7 @@
 	version = "custom"
 )
 
-func runVersion(cmd *cobra.Command, args []string) error {
+func runVersion(cmd *Command, args []string) error {
 	w := cmd.OutOrStdout()
 	fmt.Fprintf(w, "cue version %v %s/%s\n",
 		version,
diff --git a/cmd/cue/cmd/vet.go b/cmd/cue/cmd/vet.go
index 5c9818c..6e5e49a 100644
--- a/cmd/cue/cmd/vet.go
+++ b/cmd/cue/cmd/vet.go
@@ -67,12 +67,12 @@
 If more than one expression is given, all must match all values.
 `
 
-func newVetCmd() *cobra.Command {
+func newVetCmd(c *Command) *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "vet",
 		Short: "validate data",
 		Long:  vetDoc,
-		RunE:  doVet,
+		RunE:  mkRunE(c, doVet),
 	}
 
 	cmd.Flags().BoolP(string(flagConcrete), "c", false,
@@ -84,7 +84,7 @@
 	return cmd
 }
 
-func doVet(cmd *cobra.Command, args []string) error {
+func doVet(cmd *Command, args []string) error {
 	builds := loadFromArgs(cmd, args)
 	if builds == nil {
 		return nil
@@ -135,7 +135,7 @@
 	return nil
 }
 
-func vetFiles(cmd *cobra.Command, inst *cue.Instance, files []string) {
+func vetFiles(cmd *Command, inst *cue.Instance, files []string) {
 	expressions := flagExpression.StringArray(cmd)
 
 	var check cue.Value
diff --git a/cmd/cue/cmd/vet_test.go b/cmd/cue/cmd/vet_test.go
index 0786095..6856daa 100644
--- a/cmd/cue/cmd/vet_test.go
+++ b/cmd/cue/cmd/vet_test.go
@@ -17,16 +17,16 @@
 import "testing"
 
 func TestVet(t *testing.T) {
-	runCommand(t, newVetCmd(), "vet")
+	runCommand(t, newVetCmd(newRootCmd()), "vet")
 
-	cmd := newVetCmd()
+	cmd := newVetCmd(newRootCmd())
 	mustParseFlags(t, cmd, "-c")
 	runCommand(t, cmd, "vet_conc")
 
-	cmd = newVetCmd()
+	cmd = newVetCmd(newRootCmd())
 	runCommand(t, cmd, "vet_file", "./testdata/vet/vet.cue", "./testdata/vet/data.yaml")
 
-	cmd = newVetCmd()
+	cmd = newVetCmd(newRootCmd())
 	mustParseFlags(t, cmd, "-e", "File")
 	runCommand(t, cmd, "vet_expr", "./testdata/vet/vet.cue", "./testdata/vet/data.yaml")