{Core} Top-level tab completion performance improvements (#32616)

Author: DanielMicrosoft

.vscode/launch.json

@@ -38,6 +38,26 @@
                 "--help"
             ],
             "console": "integratedTerminal",
+        },
+        {
+            "name": "Azure CLI Debug Tab Completion (External Console)",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/src/azure-cli/azure/cli/__main__.py",
+            "args": [],
+            "console": "externalTerminal",
+            "cwd": "${workspaceFolder}",
+            "env": {
+                "_ARGCOMPLETE": "1",
+                "COMP_LINE": "az vm create --",
+                "COMP_POINT": "18",
+                "_ARGCOMPLETE_SUPPRESS_SPACE": "0",
+                "_ARGCOMPLETE_IFS": "\n",
+                "_ARGCOMPLETE_SHELL": "powershell",
+                "ARGCOMPLETE_USE_TEMPFILES": "1",
+                "_ARGCOMPLETE_STDOUT_FILENAME": "C:\\temp\\az_debug_completion.txt"
+            },
+            "justMyCode": false
         }
     ]
 }

src/azure-cli-core/azure/cli/core/__init__.py

@@ -28,6 +28,8 @@
 EXCLUDED_PARAMS = ['self', 'raw', 'polling', 'custom_headers', 'operation_config',
                    'content_version', 'kwargs', 'client', 'no_wait']
 EVENT_FAILED_EXTENSION_LOAD = 'MainLoader.OnFailedExtensionLoad'
+# Marker used by CommandIndex.get() to signal top-level tab completion optimization
+TOP_LEVEL_COMPLETION_MARKER = '__top_level_completion__'
 
 # [Reserved, in case of future usage]
 # Modules that will always be loaded. They don't expose commands but hook into CLI core.
@@ -225,6 +227,24 @@ def __init__(self, cli_ctx=None):
         self.cmd_to_loader_map = {}
         self.loaders = []
 
+    def _create_stub_commands_for_completion(self, command_names):
+        """Create stub commands for top-level tab completion optimization.
+
+        Stub commands allow argcomplete to parse command names without loading modules.
+
+        :param command_names: List of command names to create stubs for
+        """
+        from azure.cli.core.commands import AzCliCommand
+
+        def _stub_handler(*_args, **_kwargs):
+            """Stub command handler used only for argument completion."""
+            return None
+
+        for cmd_name in command_names:
+            if cmd_name not in self.command_table:
+                # Stub commands only need names for argcomplete parser construction.
+                self.command_table[cmd_name] = AzCliCommand(self, cmd_name, _stub_handler)
+
     def _update_command_definitions(self):
         for cmd_name in self.command_table:
             loaders = self.cmd_to_loader_map[cmd_name]
@@ -423,9 +443,16 @@ def _get_extension_suppressions(mod_loaders):
             index_result = command_index.get(args)
             if index_result:
                 index_modules, index_extensions = index_result
+
+                if index_modules == TOP_LEVEL_COMPLETION_MARKER:
+                    self._create_stub_commands_for_completion(index_extensions)
+                    _update_command_table_from_extensions([], ALWAYS_LOADED_EXTENSIONS)
+                    return self.command_table
+
                 # Always load modules and extensions, because some of them (like those in
                 # ALWAYS_LOADED_EXTENSIONS) don't expose a command, but hooks into handlers in CLI core
                 _update_command_table_from_modules(args, index_modules)
+
                 # The index won't contain suppressed extensions
                 _update_command_table_from_extensions([], index_extensions)
 
@@ -473,7 +500,6 @@ def _get_extension_suppressions(mod_loaders):
             else:
                 logger.debug("No module found from index for '%s'", args)
 
-        # No module found from the index. Load all command modules and extensions
         logger.debug("Loading all modules and extensions")
         _update_command_table_from_modules(args)
 
@@ -662,6 +688,23 @@ def __init__(self, cli_ctx=None):
             self.cloud_profile = cli_ctx.cloud.profile
         self.cli_ctx = cli_ctx
 
+    def _get_top_level_completion_commands(self):
+        """Get top-level command names for tab completion optimization.
+
+        Returns marker and list of top-level commands (e.g., 'network', 'vm') for creating
+        stub commands without module loading. Returns None if index is empty, triggering
+        fallback to full module loading.
+
+        :return: tuple of (TOP_LEVEL_COMPLETION_MARKER, list of top-level command names) or None
+        """
+        index = self.INDEX.get(self._COMMAND_INDEX) or {}
+        if not index:
+            logger.debug("Command index is empty, will fall back to loading all modules")
+            return None
+        top_level_commands = list(index.keys())
+        logger.debug("Top-level completion: %d commands available", len(top_level_commands))
+        return TOP_LEVEL_COMPLETION_MARKER, top_level_commands
+
     def get(self, args):
         """Get the corresponding module and extension list of a command.
 
@@ -681,6 +724,9 @@ def get(self, args):
         # Make sure the top-level command is provided, like `az version`.
         # Skip command index for `az` or `az --help`.
         if not args or args[0].startswith('-'):
+            # For top-level completion (az [tab])
+            if not args and self.cli_ctx.data.get('completer_active'):
+                return self._get_top_level_completion_commands()
             return None
 
         # Get the top-level command, like `network` in `network vnet create -h`

src/azure-cli-core/azure/cli/core/tests/test_argcomplete.py

@@ -64,3 +64,21 @@ def dummy_completor(*args, **kwargs):
         with open('argcomplete.out') as f:
             self.assertEqual(f.read(), 'dummystorage ')
         os.remove('argcomplete.out')
+
+    def test_top_level_completion(self):
+        """Test that top-level completion (az [tab]) returns command names from index"""
+        import os
+        import sys
+
+        if sys.platform == 'win32':
+            self.skipTest('Skip argcomplete test on Windows')
+
+        run_cmd(['az'], env=self.argcomplete_env('az ', '3'))
+        with open('argcomplete.out') as f:
+            completions = f.read().split()
+            # Verify common top-level commands are present
+            self.assertIn('account', completions)
+            self.assertIn('vm', completions)
+            self.assertIn('network', completions)
+            self.assertIn('storage', completions)
+        os.remove('argcomplete.out')