Delegate help and completion to extensions when applicable

internal/temporalcli/commands.extension.go

@@ -36,8 +36,20 @@ func (err ExtensionNonZeroExit) Unwrap() error {
 // tryExecuteExtension tries to execute an extension command if the command is not a built-in command.
 // It returns an error if the extension command fails, and a boolean indicating whether an extension was executed.
 func tryExecuteExtension(cctx *CommandContext, tcmd *TemporalCommand) (error, bool) {
+	// Special commands like "help" and "__complete" should be set aside and delegated to an extension command that matches
+	// the rest of the given arguments. "temporal help my-extension" should be rewritten to "temporal-my-extension help"
+	// Some of these commands, like "__complete" used for shell completion, don't actually get registered by Cobra until
+	// just before they're invoked, so we need to split out the delegatable commands before trying to Find() a matching
+	// subcommand.
+	delegatableCommands, nonDelegatedArgs := splitDelegatedCommands(cctx.Options.Args)
+
+	// If a completion command (used for generating the script that invokes "__complete") already exists or has been
+	// explicitly disabled, this will do nothing, but if neither of those cases are true, we want to make sure
+	// this command is registered so extensions can't shadow it.
+	tcmd.Command.InitDefaultCompletionCmd()
+
 	// Find the deepest matching built-in command and remaining args.
-	foundCmd, remainingArgs, findErr := tcmd.Command.Find(cctx.Options.Args)
+	foundCmd, remainingArgs, findErr := tcmd.Command.Find(nonDelegatedArgs)
 
 	// Cobra normally adds --help/-h before parsing, but extension dispatch
 	// pre-parses flags before Cobra's execution path runs. We Initialize it so that
@@ -58,6 +70,7 @@ func tryExecuteExtension(cctx *CommandContext, tcmd *TemporalCommand) (error, bo
 
 	// Search for an extension executable.
 	cmdPrefix := strings.Fields(foundCmd.CommandPath())
+
 	extPath, extArgs := lookupExtension(cmdPrefix, extArgs)
 
 	// Parse CLI args that need validation.
@@ -71,6 +84,16 @@ func tryExecuteExtension(cctx *CommandContext, tcmd *TemporalCommand) (error, bo
 		return nil, false
 	}
 
+	if len(delegatableCommands) > 0 && isCompletionCommand(delegatableCommands[0]) && len(extArgs) == 0 {
+		// __complete always expects at least one argument, the last of which is the current subcommand
+		// or argument to expand, with an empty string matching all possibilities.
+		// ["temporal", "__complete", "activity"] means this cli should return any subcommands and extentions that
+		// match "activity", whereas ["temporal", "__complete", "activity", ""] means we should show what's available
+		// on the activity subcommand. The same logic applies to extension commands, so even if we matched an extension,
+		// if there are no further args, it's still this cli's responsibility to respond to the completion request.
+		return nil, false
+	}
+
 	// Apply --command-timeout if set.
 	ctx := cctx.Context
 	if timeout := tcmd.CommandTimeout.Duration(); timeout > 0 {
@@ -79,7 +102,9 @@ func tryExecuteExtension(cctx *CommandContext, tcmd *TemporalCommand) (error, bo
 		defer cancel()
 	}
 
-	cmd := exec.CommandContext(ctx, extPath, append(cliPassArgs, extArgs...)...)
+	rebuiltArgs := slices.Concat(delegatableCommands, cliPassArgs, extArgs)
+
+	cmd := exec.CommandContext(ctx, extPath, rebuiltArgs...)
 	cmd.Stdin, cmd.Stdout, cmd.Stderr = cctx.Options.Stdin, cctx.Options.Stdout, cctx.Options.Stderr
 	if err := cmd.Run(); err != nil {
 		if ctx.Err() != nil {
@@ -94,6 +119,36 @@ func tryExecuteExtension(cctx *CommandContext, tcmd *TemporalCommand) (error, bo
 	return nil, true
 }
 
+// splitDelegatedCommands separates out commands that should be delegated to an extension
+// from the rest of the args given. These commands are inherently position-dependent, so they're
+// only treated specially when they're at the start of the list of arguments.
+func splitDelegatedCommands(args []string) ([]string, []string) {
+	if len(args) == 0 {
+		return args, args
+	}
+
+	if args[0] == "help" {
+		// "help __complete" never delegates, whatever comes after, so we can just mark "help" as delegatable and see what matches
+		return args[:1], args[1:]
+	}
+
+	if isCompletionCommand(args[0]) {
+		if len(args) > 1 && args[1] == "help" {
+			// "__complete help" is what happens when a user types "temporal help<TAB>", so it should delegate both. This allows
+			// shell completion to display the available help topics available from an extension
+			return args[:2], args[2:]
+		}
+
+		return args[:1], args[1:]
+	}
+
+	return []string{}, args
+}
+
+func isCompletionCommand(arg string) bool {
+	return arg == cobra.ShellCompRequestCmd || arg == cobra.ShellCompNoDescRequestCmd
+}
+
 func groupArgs(foundCmd *cobra.Command, args []string) (cliParseArgs, cliPassArgs, extArgs []string) {
 	seenPos := false
 	for i := 0; i < len(args); i++ {
@@ -199,10 +254,9 @@ func lookupExtension(cmdPrefix, extArgs []string) (string, []string) {
 }
 
 // discoverExtensions scans the PATH for executables with the "temporal-" prefix
-// and returns their command parts (without the prefix).
-func discoverExtensions() [][]string {
-	var extensions [][]string
-	seen := make(map[string]bool)
+// and returns their commands (without the prefix) mapped to the executable path
+func discoverExtensions() map[string]string {
+	extensions := make(map[string]string)
 
 	for _, dir := range filepath.SplitList(os.Getenv("PATH")) {
 		if dir == "" {
@@ -229,13 +283,12 @@ func discoverExtensions() [][]string {
 			}
 
 			path := extensionBinaryToCommandPath(baseName)
-			key := strings.Join(path, "/")
-			if seen[key] {
+			key := strings.Join(path, " ")
+			if extensions[key] != "" {
 				continue
 			}
 
-			seen[key] = true
-			extensions = append(extensions, path)
+			extensions[key] = filepath.Join(dir, entry.Name())
 		}
 	}
 	return extensions

internal/temporalcli/commands.extension_test.go

@@ -62,6 +62,30 @@ func TestExtension_PrefersMostSpecificExtension(t *testing.T) {
 	assert.Equal(t, "Args: temporal-foo-bar \n", res.Stdout.String())
 }
 
+func TestExtension_CompleteShowsExtensions(t *testing.T) {
+	h := newExtensionHarness(t)
+	h.createExtension("temporal-foo", codeEchoArgs)
+
+	res := h.Execute("__complete", "")
+	assert.Regexp(t, `foo\s+An extension command located at .*[\\/]temporal-foo`, res.Stdout.String())
+}
+
+func TestExtension_CompleteDoesNotDelegateWithoutAdditionalArgs(t *testing.T) {
+	h := newExtensionHarness(t)
+	h.createExtension("temporal-foo", codeEchoArgs)
+
+	res := h.Execute("__complete", "foo")
+	assert.Regexp(t, `foo\s+An extension command located at .*[\\/]temporal-foo`, res.Stdout.String())
+}
+
+func TestExtension_InvokesComplete(t *testing.T) {
+	h := newExtensionHarness(t)
+	h.createExtension("temporal-foo", codeEchoArgs)
+
+	res := h.Execute("__complete", "foo", "")
+	assert.Equal(t, "Args: temporal-foo __complete \n", res.Stdout.String())
+}
+
 func TestExtension_ConvertsDashToUnderscoreInLookup(t *testing.T) {
 	h := newExtensionHarness(t)
 	h.createExtension("temporal-foo-bar_baz", codeEchoArgs)
@@ -79,6 +103,7 @@ func TestExtension_DoesNotOverrideBuiltinCommand(t *testing.T) {
 	h := newExtensionHarness(t)
 	h.createExtension("temporal-workflow", codeEchoArgs)
 	h.createExtension("temporal-workflow-list", codeEchoArgs)
+	h.createExtension("temporal-completion", codeEchoArgs)
 
 	t.Run("root command", func(t *testing.T) {
 		res := h.Execute("workflow", "--help")
@@ -102,6 +127,16 @@ func TestExtension_DoesNotOverrideBuiltinCommand(t *testing.T) {
 			assert.NoError(t, res.Err)
 		}
 	})
+
+	t.Run("__complete (shell completion)", func(t *testing.T) {
+		res := h.Execute("__complete", "")
+		assert.NotContains(t, res.Stdout.String(), "temporal-workflow")
+	})
+
+	t.Run("completion script generation", func(t *testing.T) {
+		res := h.Execute("completion", "zsh")
+		assert.NotContains(t, res.Stdout.String(), "temporal-completion")
+	})
 }
 
 func TestExtension_Flags(t *testing.T) {

internal/temporalcli/commands.go

@@ -387,6 +387,12 @@ func Execute(ctx context.Context, options CommandOptions) {
 			return
 		}
 
+		if !cctx.ActuallyRanCommand && len(cctx.Options.Args) > 0 && isCompletionCommand(cctx.Options.Args[0]) {
+			// Completion was requested, but we didn't match an extension and delegate. Register all extension
+			// commands so things like "temporal cl<TABL>" will expand to "temporal cloud"
+			registerExtensionCommands(&cmd.Command)
+		}
+
 		// Run builtin command if no extension handled the command.
 		if !cctx.ActuallyRanCommand {
 			err = cmd.Command.ExecuteContext(cctx)

internal/temporalcli/commands.help.go

@@ -1,10 +1,13 @@
 package temporalcli
 
 import (
+	"cmp"
+	"fmt"
 	"slices"
 	"strings"
 
 	"github.com/spf13/cobra"
+	"golang.org/x/exp/maps"
 )
 
 // customizeHelpCommand adds the --all/-a flag to Cobra's built-in help command
@@ -53,36 +56,64 @@ func customizeHelpCommand(rootCmd *cobra.Command) {
 }
 
 // registerExtensionCommands adds discovered extensions as placeholder commands
-// so they appear in the default help output. It filters extensions based on
-// the current command's path in the hierarchy.
+// so they appear in shell completion and the default help output. It filters extensions
+// based on the current command's path in the hierarchy.
 func registerExtensionCommands(cmd *cobra.Command) {
 	cmdPath := strings.Fields(cmd.CommandPath())
-	seen := make(map[string]bool)
 
-	for _, ext := range discoverExtensions() {
+	// When built-in subcommands are nested under other subcommands (e.g. `temporal activity cancel`),
+	// they're guaranteed to have a defined parent (`temporal activity`, which has the parent `temporal`).
+	// Extension subcommands can also be nested, but when `temporal foo bar` is created via an executable
+	// named temporal-foo-bar, there's no guarantee that the `temporal foo` command exists. For the full
+	// command to show up in help and completion, placeholders must exist at every level.
+	extensionsAndExecutables := discoverExtensions()
+
+	extensionKeys := maps.Keys(extensionsAndExecutables)
+
+	// Shorter command paths first ensures the paths shown in the short description of placeholder commands
+	// point to the closest match
+	slices.SortFunc(extensionKeys, func(a, b string) int {
+		return cmp.Compare(strings.Count(a, " "), strings.Count(b, " "))
+	})
+	for _, extKey := range extensionKeys {
+		ext := strings.Split(extKey, " ")
 		// Extension must be deeper than current command and share the same prefix
 		if len(ext) <= len(cmdPath) || !slices.Equal(ext[:len(cmdPath)], cmdPath) {
 			continue
 		}
 
-		// Get the next level command name
-		nextPart := ext[len(cmdPath)]
+		extPath := ext[len(cmdPath):]
 
-		// Skip if already added
-		if seen[nextPart] {
-			continue
-		}
+		parent := cmd
+		executablePath := extensionsAndExecutables[extKey]
 
-		// Skip if a built-in command exists
-		if found, _, _ := cmd.Find([]string{nextPart}); found != cmd {
-			continue
+		for i, nextPart := range extPath {
+			if found, _, _ := parent.Find([]string{nextPart}); found != parent {
+				// Because we order extensions by depth, we can trust that any command that already exists already
+				// has the most correct definition for its depth.
+				parent = found
+				continue
+			}
+
+			var short string
+
+			if i == len(extPath)-1 {
+				short = fmt.Sprintf("An extension command located at %s", executablePath)
+			} else {
+				short = fmt.Sprintf("Extension commands under %s", strings.Join(ext[:len(cmdPath)+i+1], " "))
+			}
+
+			newCmd := &cobra.Command{
+				Use: nextPart,
+				// Short descriptions must be unique at a given level because otherwise shell completion will
+				// group all commands with the same description on the same line
+				Short:              short,
+				DisableFlagParsing: true,
+				Run:                func(*cobra.Command, []string) {},
+			}
+			parent.AddCommand(newCmd)
+			parent = newCmd
 		}
 
-		seen[nextPart] = true
-		cmd.AddCommand(&cobra.Command{
-			Use:                nextPart,
-			DisableFlagParsing: true,
-			Run:                func(*cobra.Command, []string) {},
-		})
 	}
 }