feat(completion): improve @N<TAB> UX with grouped section menu

cli/cluster.go

@@ -61,6 +61,42 @@ func completionClusterFilter() map[uint32]bool {
 	return targetEndpointClusterIDs
 }
 
+// topLevelCommandsForCompletion snapshots the visible root subcommands and
+// labels each with its completion group ("device", "cluster", or "tool") and
+// whether it remains relevant once a node target has been selected. The
+// completion package uses the result to populate "@N+<cmd>" expansion tokens
+// so that Tab after an exact node match offers device commands, tools, and
+// help alongside endpoint completions.
+//
+// Shorthand cluster commands and hidden commands are omitted — they are
+// surfaced elsewhere (by cluster-name matching or via the cluster parent
+// command) and would otherwise flood the menu.
+func topLevelCommandsForCompletion() []completion.TopLevelCommand {
+	cmds := allRootCommands()
+	out := make([]completion.TopLevelCommand, 0, len(cmds))
+	for _, c := range cmds {
+		if c.Hidden || isShorthandCluster(c) || !c.IsAvailableCommand() {
+			continue
+		}
+		var group string
+		switch c.GroupID {
+		case groupDevices:
+			group = "device"
+		case groupClusters:
+			group = "cluster"
+		case groupTools:
+			group = "tool"
+		}
+		out = append(out, completion.TopLevelCommand{
+			Name:        c.Name(),
+			Short:       c.Short,
+			Group:       group,
+			TargetAware: !targetUnawareCommands[c.Name()],
+		})
+	}
+	return out
+}
+
 func init() {
 	allRootCommands = func() []*cobra.Command { return rootCmd.Commands() }
 	rootCmd.AddCommand(withGroup(newClusterCmd(), groupClusters))

cli/completion.go

@@ -497,7 +497,8 @@ compdef _matter matter
 # Group headers: requires group-name '' to be set (oh-my-zsh sets this
 # globally; we set it locally for matter so vanilla zsh also benefits).
 zstyle ':completion:*:matter:*' group-name ''
-zstyle ':completion:*:matter:*:targets'            format $'\e[35m── Targets ──\e[0m'
+zstyle ':completion:*:matter:*:devices'            format $'\e[35m── Devices ──\e[0m'
+zstyle ':completion:*:matter:*:endpoints'          format $'\e[35m── Endpoints ──\e[0m'
 zstyle ':completion:*:matter:*:device-commands'    format $'\e[36m── Device Commands ──\e[0m'
 zstyle ':completion:*:matter:*:cluster-commands'   format $'\e[32m── Cluster Commands ──\e[0m'
 zstyle ':completion:*:matter:*:tools'              format $'\e[33m── Tools ──\e[0m'
@@ -513,8 +514,11 @@ _matter_group_map=(
 
 _matter() {
   local -a request_cmd
-  local -a device_cmds cluster_cmds tool_cmds target_cmds other_cmds invoke_cmds attr_cmds
+  local -a device_targets endpoint_targets
+  local -a device_cmds cluster_cmds tool_cmds other_cmds invoke_cmds attr_cmds
+  local -a exp_device_cmds exp_cluster_cmds exp_tool_cmds
   local out word desc entry tag i directive _matter_line
+  local exp_target=""
 
   # Build the __complete call from the current word list.
   request_cmd=("${words[1]}" "__complete")
@@ -523,7 +527,10 @@ _matter() {
   done
   request_cmd+=("${words[$CURRENT]}")
 
-  out=$("${request_cmd[@]}" 2>/dev/null)
+  # MATTER_COMPLETION_EXPAND=zsh opts this shell into the "@N+<cmd>" expansion
+  # tokens the loop below parses. Other shells (bash/fish/powershell) do not
+  # set this and therefore never see the zsh-specific encoding.
+  out=$(MATTER_COMPLETION_EXPAND=zsh "${request_cmd[@]}" 2>/dev/null)
 
   # Extract the cobra ShellCompDirective from the trailing :N line so we can
   # honour flags like ShellCompDirectiveNoSpace (bit 1, value 2).
@@ -537,8 +544,28 @@ _matter() {
     [[ -z "$word" || "$word" == :* || "$word" == _activeHelp_* ]] && continue
     # Escape colons in word and description (zsh _describe uses : as separator).
     entry="${word//:/\\:}:${desc//:/\\:}"
-    if [[ "$word" == @* ]]; then
-      target_cmds+=("$entry")
+    if [[ "$word" == @*+* ]]; then
+      # Expansion token "@N+<cmd>": splits into the @N target prefix and a
+      # bare subcommand name. The display is the subcommand alone; selection
+      # inserts "@N <cmd>" (as two shell words) via compadd -U -p below.
+      local _exp_prefix="${word%%+*}"
+      local _exp_cmd="${word#*+}"
+      exp_target="$_exp_prefix"
+      tag="${_matter_group_map[$_exp_cmd]}"
+      local _exp_entry="${_exp_cmd//:/\\:}:${desc//:/\\:}"
+      case "$tag" in
+        device)  exp_device_cmds+=("$_exp_entry") ;;
+        cluster) exp_cluster_cmds+=("$_exp_entry") ;;
+        tool)    exp_tool_cmds+=("$_exp_entry") ;;
+        # Ungrouped commands (e.g. "help", which cobra does not register in
+        # any group) are routed to the Tools bucket so they remain visible
+        # instead of being silently dropped.
+        *)       exp_tool_cmds+=("$_exp_entry") ;;
+      esac
+    elif [[ "$word" == @*/* ]]; then
+      endpoint_targets+=("$entry")
+    elif [[ "$word" == @* ]]; then
+      device_targets+=("$entry")
     else
       tag="${_matter_group_map[$word]}"
       case "$tag" in
@@ -566,7 +593,22 @@ _matter() {
   local -a nospace
   (( directive & 2 )) && nospace=(-S '')
 
-  (( ${#target_cmds}  )) && _describe -t targets            "Targets"                target_cmds  "${nospace[@]}"
+  (( ${#device_targets}   )) && _describe -t devices            "Devices"                device_targets   "${nospace[@]}"
+  (( ${#endpoint_targets} )) && _describe -t endpoints          "Endpoints"              endpoint_targets "${nospace[@]}"
+
+  # Expansion entries share the same "@N " target prefix. -U disables prefix
+  # matching against the typed @N word so the bare subcommand names are kept
+  # as candidates; -p prepends "@N " to the inserted text so the subcommand
+  # becomes a separate shell word after the target. -Q suppresses zsh's
+  # default quoting so the space in "@N " stays a plain word separator
+  # rather than being inserted as a backslash-escaped "\ ".
+  if [[ -n "$exp_target" ]]; then
+    local -a exp_prefix_arg=(-U -Q -p "${exp_target} ")
+    (( ${#exp_device_cmds}  )) && _describe -t device-commands    "Device Commands"   exp_device_cmds  "${exp_prefix_arg[@]}"
+    (( ${#exp_cluster_cmds} )) && _describe -t cluster-commands   "Cluster Commands"  exp_cluster_cmds "${exp_prefix_arg[@]}"
+    (( ${#exp_tool_cmds}    )) && _describe -t tools              "Tools"             exp_tool_cmds    "${exp_prefix_arg[@]}"
+  fi
+
   (( ${#device_cmds}  )) && _describe -t device-commands    "Device Commands"        device_cmds  "${nospace[@]}"
   (( ${#cluster_cmds} )) && _describe -t cluster-commands   "Cluster Commands"       cluster_cmds "${nospace[@]}"
   (( ${#tool_cmds}    )) && _describe -t tools              "Tools"                  tool_cmds    "${nospace[@]}"

cli/completion/completer.go

@@ -9,6 +9,7 @@ package completion
 
 import (
 	"fmt"
+	"os"
 	"sort"
 	"strings"
 	"time"
@@ -263,10 +264,36 @@ func NodeIDCompletionFunc() func(cmd *cobra.Command, args []string, toComplete s
 	}
 }
 
+// TopLevelCommand describes a single root-level cobra subcommand for the
+// purposes of @target completion expansion. When the user types "@N<TAB>"
+// against an existing node, these commands are surfaced as "@N+<name>"
+// tokens so the shell can offer them alongside endpoint and sibling-node
+// completions in a single menu.
+type TopLevelCommand struct {
+	// Name is the subcommand name as registered with cobra (e.g. "tree").
+	Name string
+	// Short is the one-line description shown next to the command.
+	Short string
+	// Group is "device", "cluster", "tool", or "". It is metadata for
+	// shell-specific completion helpers and is not encoded in the emitted
+	// "@N+<name>" token; the zsh script derives grouping from its own
+	// statically-generated _matter_group_map, keyed by command name.
+	Group string
+	// TargetAware reports whether the command accepts/requires a node
+	// target. Commands that operate on the fabric as a whole (e.g.
+	// "commission", "discover") are skipped for @N expansion so the menu
+	// stays relevant.
+	TargetAware bool
+}
+
 // RootCompletionFunc returns a cobra ValidArgsFunction for the root command
-// that handles two completion types:
+// that handles three completion types:
 //
 //   - @target tokens (e.g. "@1/2") — delegated to TargetCompletionFunc.
+//     When a numeric @N exactly matches an existing node, the returned set
+//     is enriched with endpoint tokens (@N/0, @N/1, ...) and "@N+<cmd>"
+//     expansion tokens for the commands returned by topLevelCommands, so
+//     the user does not need to type " " or "/" to see what comes next.
 //   - cluster shorthand commands — case-insensitive prefix/substring match of
 //     cluster names, so typing "on<TAB>" offers "OnOff" and "level<TAB>" offers
 //     "LevelControl".
@@ -278,11 +305,15 @@ func NodeIDCompletionFunc() func(cmd *cobra.Command, args []string, toComplete s
 // to the set of cluster IDs present on the current target endpoint. A nil map
 // means no filter (show all clusters); a non-nil but empty map means no
 // clusters are applicable (e.g. node-only target without an endpoint).
+//
+// topLevelCommands, if non-nil, is called when expanding an exact @N match to
+// seed the "Device Commands" / "Tools" / "Cluster Commands" sections.
 func RootCompletionFunc(
 	registry *clusters.Registry,
 	allowedClusters func() map[uint32]bool,
+	topLevelCommands func() []TopLevelCommand,
 ) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-	targetFn := TargetCompletionFunc()
+	targetFn := TargetCompletionFunc(topLevelCommands)
 	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		if strings.HasPrefix(toComplete, "@") {
 			return targetFn(cmd, args, toComplete)
@@ -314,6 +345,21 @@ func RootCompletionFunc(
 	}
 }
 
+// ExpandSeparator is the separator embedded in "@N+<cmd>" expansion tokens.
+// The zsh completion script recognises it to split the target prefix from the
+// subcommand name and invoke compadd with -U -p "@N " so the subcommand is
+// inserted as a separate shell word while the displayed candidate is the
+// bare subcommand name. The "+" character is used because it cannot appear
+// in either a valid numeric node token or a top-level cobra command name.
+const ExpandSeparator = "+"
+
+// ExpandEnvVar is the environment variable the zsh completion script sets
+// (to "zsh") before invoking "matter __complete ...". Its presence opts the
+// caller into receiving "@N+<cmd>" expansion tokens. Other shells (bash,
+// fish, PowerShell) do not set it, so they continue to receive only plain
+// @N / @N/<ep> tokens and never display the literal expansion syntax.
+const ExpandEnvVar = "MATTER_COMPLETION_EXPAND"
+
 // TargetCompletionFunc returns a cobra ValidArgsFunction that completes
 // @target tokens in two stages. Emitted tokens are always numeric @N;
 // device names are shown only in the description as a visual hint.
@@ -327,11 +373,24 @@ func RootCompletionFunc(
 // "/" to proceed to endpoint selection or " " (space) for device
 // commands.
 //
+// Stage 1b — exact match expansion:
+// When toComplete is a numeric @N that exactly identifies a commissioned
+// node, the result set is enriched with:
+//   - endpoint tokens "@N/0", "@N/1", ... so the user sees endpoints without
+//     having to type "/" first; and
+//   - "@N+<cmd>" expansion tokens for each top-level command returned by
+//     topLevelCommands (if non-nil), so device-level commands and tools
+//     appear in the same menu as the endpoint list. The zsh script strips
+//     the "@N+" prefix before display and inserts the selection as a
+//     separate word after "@N ".
+//
 // Stage 2 — endpoint selection ("/" present):
 // When the user types "@1/", completions are the non-root endpoints on
 // the matched node (e.g. "@1/1", "@1/2"). Normal trailing space is
 // applied so the user can proceed to a command after selection.
-func TargetCompletionFunc() func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+func TargetCompletionFunc(
+	topLevelCommands func() []TopLevelCommand,
+) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		// Only complete if the user is typing a @target token.
 		if !strings.HasPrefix(toComplete, "@") {
@@ -356,6 +415,18 @@ func TargetCompletionFunc() func(cmd *cobra.Command, args []string, toComplete s
 		partial := strings.ToLower(toComplete[1:])
 		namePart, epPart, hasSlash := strings.Cut(partial, "/")
 
+		// Look for an exact numeric match so we can enrich stage-1 with
+		// endpoint + subcommand expansion for that node.
+		var exactNode *store.Node
+		if !hasSlash && isAllDigits(namePart) {
+			for _, n := range nodes {
+				if fmt.Sprintf("%d", n.ID) == namePart {
+					exactNode = n
+					break
+				}
+			}
+		}
+
 		var completions []string
 		for _, n := range nodes {
 			idStr := fmt.Sprintf("%d", n.ID)
@@ -419,9 +490,51 @@ func TargetCompletionFunc() func(cmd *cobra.Command, args []string, toComplete s
 			}
 		}
 
+		// ── Stage 1b: enrich exact-match with endpoints + command expansion ──
+		if exactNode != nil {
+			idStr := fmt.Sprintf("%d", exactNode.ID)
+			alias := idStr
+			if exactNode.Name != "" {
+				alias = strings.ReplaceAll(strings.ToLower(exactNode.Name), " ", "-")
+			}
+
+			// Endpoint tokens. User sees these in the same menu; they match
+			// the "@N" prefix so the shell keeps them.
+			for _, ep := range exactNode.Endpoints {
+				epStr := fmt.Sprintf("%d", ep.ID)
+				epDesc := endpointDescription(ep)
+				var desc string
+				if alias != idStr {
+					desc = fmt.Sprintf("%s  [%s]", epDesc, alias)
+				} else {
+					desc = epDesc
+				}
+				completions = append(completions, fmt.Sprintf("@%s/%s\t%s", idStr, epStr, desc))
+			}
+
+			// Command expansion tokens. The zsh script splits on
+			// ExpandSeparator and inserts "@N " as a word prefix via
+			// compadd -U -p so the subcommand is displayed bare but
+			// dispatched as a separate shell word. Only emitted when the
+			// caller opted in via ExpandEnvVar: bash, fish, and PowerShell
+			// completions do not rewrite these tokens and would otherwise
+			// surface the literal "@N+<cmd>" text to the user.
+			if topLevelCommands != nil && os.Getenv(ExpandEnvVar) == "zsh" {
+				for _, tc := range topLevelCommands() {
+					if !tc.TargetAware {
+						continue
+					}
+					completions = append(completions,
+						fmt.Sprintf("@%s%s%s\t%s", idStr, ExpandSeparator, tc.Name, tc.Short))
+				}
+			}
+		}
+
 		if !hasSlash {
 			// Node-only stage: suppress trailing space so the user can type
-			// "/" for endpoint or " " for device-level commands.
+			// "/" for endpoint or " " for device-level commands. The shell
+			// script intercepts "@N+<cmd>" tokens separately and inserts
+			// them with a literal space via compadd -U -p.
 			return completions, cobra.ShellCompDirectiveNoSpace | cobra.ShellCompDirectiveNoFileComp
 		}
 		return completions, cobra.ShellCompDirectiveNoFileComp