Fix Invalidate command index after azdev extension add/remove (#543)

* fix: invalidate command index after azdev extension add/remove

* fix: ci error

* fix: ci error

* fix: resolve conflicts

* Update HISTORY.rst

Co-authored-by: necusjz <necusjz@gmail.com>

* Update HISTORY.rst

* Update HISTORY.rst

---------

Co-authored-by: necusjz <necusjz@gmail.com>

Author: huiii99

HISTORY.rst

@@ -2,6 +2,10 @@
 
 Release History
 ===============
+0.2.11
+++++++
+* `azdev extension add/remove`: Invalidate command index after installing or removing extensions.
+
 0.2.10
 ++++++
 * Add support for Python 3.14 and drop support for Python 3.9

azdev/operations/extensions/__init__.py

@@ -16,11 +16,36 @@
 
 from azdev.utilities import (
     cmd, py_cmd, pip_cmd, display, get_ext_repo_paths, find_files, get_azure_config, get_azdev_config,
-    require_azure_cli, heading, subheading, EXTENSION_PREFIX)
+    get_azure_config_dir, require_azure_cli, heading, subheading, EXTENSION_PREFIX)
 from .version_upgrade import VersionUpgradeMod
 
 logger = get_logger(__name__)
 
+# These are the index files cleared by CommandIndex().invalidate() in azure-cli-core.
+# Refer: azure-cli-core/azure/cli/core/__init__.py
+_COMMAND_INDEX_FILES = (
+    'commandIndex.json',
+    'extensionIndex.json',
+    'helpIndex.json',
+    'extensionHelpIndex.json',
+)
+
+
+def _invalidate_command_index():
+    """Delete the CLI command index files so they are regenerated on next invocation.
+
+    This mirrors the behavior of ``CommandIndex().invalidate()`` in azure-cli-core but
+    works without requiring a fully initialized CLI session.
+    """
+    azure_config_dir = get_azure_config_dir()
+    for filename in _COMMAND_INDEX_FILES:
+        filepath = os.path.join(azure_config_dir, filename)
+        try:
+            os.remove(filepath)
+            logger.debug("Deleted command index file: %s", filepath)
+        except OSError:
+            pass
+
 
 def add_extension(extensions):
 
@@ -49,6 +74,8 @@ def add_extension(extensions):
         if result.error:
             raise result.error  # pylint: disable=raising-bad-type
 
+    _invalidate_command_index()
+
 
 def remove_extension(extensions):
 
@@ -83,6 +110,8 @@ def remove_extension(extensions):
                 display("Removing '{}'...".format(path_to_remove))
                 shutil.rmtree(path_to_remove)
 
+    _invalidate_command_index()
+
 
 def _get_installed_dev_extensions(dev_sources):
     from glob import glob

azdev/operations/tests/test_extension_index_invalidation.py

@@ -0,0 +1,81 @@
+# -----------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# -----------------------------------------------------------------------------
+
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from azdev.operations.extensions import _invalidate_command_index, _COMMAND_INDEX_FILES
+
+
+class TestInvalidateCommandIndex(unittest.TestCase):
+
+    def test_deletes_all_index_files(self):
+        """All four command index files should be deleted when they exist."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            for filename in _COMMAND_INDEX_FILES:
+                filepath = os.path.join(tmpdir, filename)
+                with open(filepath, 'w') as f:
+                    f.write('{}')
+
+            with patch('azdev.operations.extensions.get_azure_config_dir', return_value=tmpdir):
+                _invalidate_command_index()
+
+            for filename in _COMMAND_INDEX_FILES:
+                self.assertFalse(os.path.exists(os.path.join(tmpdir, filename)),
+                                 f'{filename} should have been deleted')
+
+    def test_handles_missing_files_gracefully(self):
+        """Should not raise when index files do not exist."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            with patch('azdev.operations.extensions.get_azure_config_dir', return_value=tmpdir):
+                _invalidate_command_index()  # should not raise
+
+    def test_handles_partial_files(self):
+        """Should delete existing files and skip missing ones without error."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            existing = _COMMAND_INDEX_FILES[0]
+            with open(os.path.join(tmpdir, existing), 'w') as f:
+                f.write('{}')
+
+            with patch('azdev.operations.extensions.get_azure_config_dir', return_value=tmpdir):
+                _invalidate_command_index()
+
+            self.assertFalse(os.path.exists(os.path.join(tmpdir, existing)))
+
+    @patch('azdev.operations.extensions._invalidate_command_index')
+    @patch('azdev.operations.extensions.pip_cmd')
+    @patch('azdev.operations.extensions.find_files', return_value=['/repo/src/my-ext/setup.py'])
+    @patch('azdev.operations.extensions.get_ext_repo_paths', return_value=['/repo'])
+    def test_add_extension_calls_invalidate(self, _mock_paths, _mock_find, mock_pip, mock_invalidate):
+        """add_extension should call _invalidate_command_index after installing."""
+        from unittest.mock import MagicMock
+        mock_pip.return_value = MagicMock(error=None)
+
+        from azdev.operations.extensions import add_extension
+        add_extension(['my-ext'])
+
+        mock_invalidate.assert_called_once()
+
+    @patch('azdev.operations.extensions._invalidate_command_index')
+    @patch('azdev.operations.extensions.pip_cmd')
+    @patch('azdev.operations.extensions.display')
+    @patch('azdev.operations.extensions.find_files', return_value=['/repo/src/my-ext/my_ext.egg-info'])
+    @patch('azdev.operations.extensions.get_ext_repo_paths', return_value=['/repo'])
+    def test_remove_extension_calls_invalidate(
+            self, _mock_paths, _mock_find, _mock_display,
+            _mock_pip, mock_invalidate):
+        """remove_extension should call _invalidate_command_index after removing."""
+        with patch('os.listdir', return_value=[]):
+            from azdev.operations.extensions import remove_extension
+            remove_extension(['my-ext'])
+
+        mock_invalidate.assert_called_once()
+
+
+if __name__ == '__main__':
+    unittest.main()