cobra: fix registering prediction for persistent flag

cmpcobra/completer.go

@@ -1,6 +1,7 @@
 package cmpcobra
 
 import (
+	"cmp"
 	"fmt"
 	"strings"
 
@@ -26,7 +27,10 @@ var (
 // command "factories" to share that value around. The registered values are unique
 // pointers so overwriting won't happen.
 func RegisterFlag(cmd *cobra.Command, name string, predictor predict.Predictor) {
-	flag := cmd.Flags().Lookup(name)
+	local := cmd.Flags().Lookup(name)
+	persistent := cmd.PersistentFlags().Lookup(name)
+
+	flag := cmp.Or(local, persistent)
 	if flag == nil {
 		panic(fmt.Sprintf("flag %q doesn't exist on %q", name, cmd.Name()))
 	}

cmpcobra/completer_test.go

@@ -47,6 +47,19 @@ func TestQuery(t *testing.T) {
 	cmptest.Assert(t, New(cmd), "query -c <TAB> -c colX table2", []string{})
 }
 
+func TestPersistentRegister(t *testing.T) {
+	root := &cobra.Command{
+		Use: "root",
+	}
+	root.PersistentFlags().String("env", "", "")
+	pred := predict.Func(func(args args.Args) []string {
+		return []string{"foo"}
+	})
+	RegisterFlag(root, "env", pred)
+
+	cmptest.Assert(t, New(root), "root --env <TAB>", []string{"foo"})
+}
+
 func TestPersistent(t *testing.T) {
 	// Test that the child command can get access to persistent flags without looking
 	// at the root