fix: fix top-level help and fix help cache write conditions

Author: DanielMicrosoft

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

@@ -476,6 +476,14 @@ def _get_extension_suppressions(mod_loaders):
                 # The index won't contain suppressed extensions
                 _update_command_table_from_extensions([], index_extensions)
 
+                # If we loaded all extensions as a safety fallback, refresh extension overlay cache
+                # so subsequent runs can use blended targeted loading.
+                if use_command_index and index_extensions is None and command_index.cloud_profile == 'latest':
+                    command_index.update_extension_index(self.command_table)
+                    # We already paid the cost to load all extensions. Refresh help overlay as well so
+                    # top-level help can stay on packaged base + extension overlay fast path.
+                    self._cache_help_index(command_index)
+
                 logger.debug("Loaded %d groups, %d commands.", len(self.command_group_table), len(self.command_table))
                 from azure.cli.core.util import roughly_parse_command
                 # The index may be outdated. Make sure the command appears in the loaded command table
@@ -517,6 +525,15 @@ def _get_extension_suppressions(mod_loaders):
 
                 logger.debug("Could not find a match in the command or command group table for '%s'. "
                              "The index may be outdated.", raw_cmd)
+
+                if command_index.cloud_profile == 'latest' and lookup_args and \
+                        not self.cli_ctx.data['completer_active']:
+                    top_command = lookup_args[0]
+                    packaged_core_index = command_index._get_packaged_command_index(ignore_extensions=True) or {}
+                    if top_command != 'help' and top_command not in packaged_core_index:
+                        logger.debug("Top-level command '%s' is not in packaged core index. "
+                                     "Skipping full core module reload.", top_command)
+                        return self.command_table
             else:
                 logger.debug("No module found from index for '%s'", args)
 
@@ -1048,12 +1065,19 @@ def get(self, args):
                 if result:
                     return result
 
-                if force_load_all_extensions and normalized_args and not normalized_args[0].startswith('-') and \
-                        not self.cli_ctx.data['completer_active']:
-                    logger.debug("No match found in blended latest index for '%s'. Loading all extensions.",
-                                 normalized_args[0])
-                    # Load all extensions to resolve extension-only top-level commands without rebuilding all modules.
-                    return [], None
+                if normalized_args and not normalized_args[0].startswith('-') and \
+                    not self.cli_ctx.data['completer_active'] and not force_packaged_for_version and \
+                    top_command != 'help':
+                    # Unknown top-level command on latest should prefer extension-only retry and avoid
+                    # full core module rebuild to preserve packaged-index startup benefit.
+                    if has_non_always_loaded_extensions:
+                        logger.debug("No match found in blended latest index for '%s'. Loading all extensions.",
+                                     normalized_args[0])
+                        return [], None
+
+                    logger.debug("No match found in latest index for '%s' and no dynamic extensions are installed. "
+                                 "Skipping core module rebuild.", normalized_args[0])
+                    return [], []
 
                 logger.debug("No match found in blended latest index. Falling back to local command index.")
 
@@ -1225,21 +1249,27 @@ def update(self, command_table):
         elapsed_time = timeit.default_timer() - start_time
         self.INDEX[self._COMMAND_INDEX] = index
 
-        # Maintain extension-only overlay for latest profile so packaged core can be blended.
-        if self.cloud_profile == 'latest':
-            extension_index = defaultdict(list)
-            for command_name, command in command_table.items():
-                top_command = command_name.split()[0]
-                module_name = command.loader.__module__
-                if module_name.startswith('azext_') and module_name not in extension_index[top_command]:
-                    extension_index[top_command].append(module_name)
-
-            self.EXTENSION_INDEX[self._COMMAND_INDEX_VERSION] = __version__
-            self.EXTENSION_INDEX[self._COMMAND_INDEX_CLOUD_PROFILE] = self.cloud_profile
-            self.EXTENSION_INDEX[self._COMMAND_INDEX] = extension_index
+        self.update_extension_index(command_table)
 
         logger.debug("Updated command index in %.3f seconds.", elapsed_time)
 
+    def update_extension_index(self, command_table):
+        """Update extension-only overlay index from a command table (latest profile only)."""
+        if self.cloud_profile != 'latest':
+            return
+
+        from collections import defaultdict
+        extension_index = defaultdict(list)
+        for command_name, command in command_table.items():
+            top_command = command_name.split()[0]
+            module_name = command.loader.__module__
+            if module_name.startswith('azext_') and module_name not in extension_index[top_command]:
+                extension_index[top_command].append(module_name)
+
+        self.EXTENSION_INDEX[self._COMMAND_INDEX_VERSION] = __version__
+        self.EXTENSION_INDEX[self._COMMAND_INDEX_CLOUD_PROFILE] = self.cloud_profile
+        self.EXTENSION_INDEX[self._COMMAND_INDEX] = extension_index
+
     def invalidate(self):
         """Invalidate the command index.
 

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

@@ -535,7 +535,7 @@ def test_command_index_handles_leading_output_option(self):
     @mock.patch('azure.cli.core.extension.get_extension_modname', _mock_get_extension_modname)
     @mock.patch('azure.cli.core.extension.get_extensions', _mock_get_extensions)
     def test_command_index_loads_all_extensions_when_overlay_missing(self):
-        from azure.cli.core._session import INDEX
+        from azure.cli.core._session import INDEX, EXTENSION_INDEX, EXTENSION_HELP_INDEX
         from azure.cli.core import CommandIndex, __version__
 
         cli = DummyCli()
@@ -560,6 +560,54 @@ def test_command_index_loads_all_extensions_when_overlay_missing(self):
         # Missing overlay triggers loading all extensions, but avoids full module rebuild.
         self.assertEqual(list(cmd_tbl), ['hello mod-only', 'hello overridden', 'hello ext-only'])
         self.assertEqual(INDEX[CommandIndex._COMMAND_INDEX], {})
+        self.assertEqual(EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_VERSION], __version__)
+        self.assertEqual(EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_CLOUD_PROFILE], cli.cloud.profile)
+        self.assertIn('hello', EXTENSION_INDEX[CommandIndex._COMMAND_INDEX])
+        self.assertIn('azext_hello1', EXTENSION_INDEX[CommandIndex._COMMAND_INDEX]['hello'])
+        self.assertIn('azext_hello2', EXTENSION_INDEX[CommandIndex._COMMAND_INDEX]['hello'])
+        self.assertEqual(EXTENSION_HELP_INDEX[CommandIndex._COMMAND_INDEX_VERSION], __version__)
+        self.assertEqual(EXTENSION_HELP_INDEX[CommandIndex._COMMAND_INDEX_CLOUD_PROFILE], cli.cloud.profile)
+        self.assertIn('groups', EXTENSION_HELP_INDEX[CommandIndex._HELP_INDEX])
+        self.assertIn('commands', EXTENSION_HELP_INDEX[CommandIndex._HELP_INDEX])
+
+    @mock.patch('importlib.import_module', _mock_import_lib)
+    @mock.patch('pkgutil.iter_modules', _mock_iter_modules)
+    @mock.patch('azure.cli.core.commands._load_command_loader', _mock_load_command_loader)
+    @mock.patch('azure.cli.core.extension.get_extension_modname', _mock_get_extension_modname)
+    @mock.patch('azure.cli.core.extension.get_extensions', _mock_get_extensions)
+    def test_command_index_latest_unknown_non_core_skips_full_core_reload(self):
+        from azure.cli.core._session import INDEX, EXTENSION_INDEX
+        from azure.cli.core import CommandIndex, __version__
+
+        cli = DummyCli()
+        loader = cli.commands_loader
+
+        INDEX[CommandIndex._COMMAND_INDEX_VERSION] = ""
+        INDEX[CommandIndex._COMMAND_INDEX_CLOUD_PROFILE] = ""
+        INDEX[CommandIndex._COMMAND_INDEX] = {}
+        EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_VERSION] = ""
+        EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_CLOUD_PROFILE] = ""
+        EXTENSION_INDEX[CommandIndex._COMMAND_INDEX] = {}
+
+        packaged_index = {
+            CommandIndex._COMMAND_INDEX_VERSION: __version__,
+            CommandIndex._COMMAND_INDEX_CLOUD_PROFILE: cli.cloud.profile,
+            CommandIndex._COMMAND_INDEX: {
+                'hello': ['azure.cli.command_modules.hello'],
+                'extra': ['azure.cli.command_modules.extra']
+            }
+        }
+
+        with mock.patch.object(CommandIndex, '_load_packaged_command_index', return_value=packaged_index):
+            cmd_tbl = loader.load_command_table(["foobar", "list"])
+
+        # Unknown non-core top-level command should try extensions without rebuilding all core modules.
+        self.assertNotIn('hello mod-only', cmd_tbl)
+        self.assertNotIn('extra final', cmd_tbl)
+        self.assertEqual(INDEX[CommandIndex._COMMAND_INDEX], {})
+        self.assertEqual(EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_VERSION], __version__)
+        self.assertEqual(EXTENSION_INDEX[CommandIndex._COMMAND_INDEX_CLOUD_PROFILE], cli.cloud.profile)
+        self.assertIn('hello', EXTENSION_INDEX[CommandIndex._COMMAND_INDEX])
 
     @mock.patch('importlib.import_module', _mock_import_lib)
     @mock.patch('pkgutil.iter_modules', _mock_iter_modules)