fix(extensions): preserve keep-config leftover on reinstall (data loss)

src/specify_cli/extensions/__init__.py

@@ -1312,6 +1312,44 @@ def check_compatibility(
 
         return True
 
+    @staticmethod
+    def _atomic_restore_config(
+        dest: Path, data: bytes, mode: int, atime: float, mtime: float
+    ) -> None:
+        """Restore a preserved config to *dest* atomically (temp file + os.replace).
+
+        Writes to a sibling temp file, then ``os.replace()`` swaps it into place.
+        ``os.replace`` operates on the directory entry, so it (a) never follows a
+        symlink that may occupy *dest* — closing the check-then-use race where a
+        racing process swaps in a symlink to overwrite an external target — and
+        (b) leaves either the old file or the fully-written new one, never a
+        partial write. Mirrors ``shared_infra._write_shared_bytes``. Mode and
+        timestamps are set on the temp file (independently — either can succeed
+        on its own) so the installed file carries them from the first moment.
+        """
+        # os.replace cannot replace a directory with a file; clear a stray
+        # (non-symlink) directory at the path first. A symlink is handled by
+        # os.replace itself (it swaps the link entry without following it).
+        if dest.is_dir() and not dest.is_symlink():
+            shutil.rmtree(dest, ignore_errors=True)
+        fd, temp_name = tempfile.mkstemp(prefix=f".{dest.name}.", dir=dest.parent)
+        temp_path = Path(temp_name)
+        try:
+            with os.fdopen(fd, "wb") as fh:
+                fh.write(data)
+            try:
+                temp_path.chmod(mode & 0o7777)  # permission bits only
+            except OSError:
+                pass
+            try:
+                os.utime(temp_path, (atime, mtime))
+            except OSError:
+                pass
+            os.replace(temp_path, dest)
+        finally:
+            if temp_path.exists():
+                temp_path.unlink()
+
     def install_from_directory(
         self,
         source_dir: Path,
@@ -1397,12 +1435,89 @@ def install_from_directory(
                 backup_config_dir.unlink()
             did_remove = self.remove(manifest.id)
 
-        # Install extension (dest_dir computed above during self-install guard)
+        # A prior `remove(keep_config=True)` leaves the top-level
+        # *-config.yml / *-config.local.yml in place while dropping the
+        # registry entry. A later reinstall is not a --force removal
+        # (did_remove is False), so the rmtree below would wipe those
+        # preserved configs and the backup-restore path never runs. Capture
+        # them here and write them back after the fresh copytree, mirroring
+        # the *-config filter the --force restore path uses.
+        # Capture (bytes, mode, atime, mtime) so restore preserves both
+        # permissions and timestamps — truly mirroring the --force restore's
+        # shutil.copy2 (which preserves mode/mtime). Config files may hold
+        # secrets (API keys); recreating them with default perms could widen
+        # access.
+        preserved_configs: dict[str, tuple[bytes, int, float, float]] = {}
         if dest_dir.exists():
-            shutil.rmtree(dest_dir)
+            for cfg_file in dest_dir.iterdir():
+                if (
+                    cfg_file.is_file()
+                    and not cfg_file.is_symlink()
+                    and (
+                        cfg_file.name.endswith("-config.yml")
+                        or cfg_file.name.endswith("-config.local.yml")
+                    )
+                ):
+                    st = cfg_file.stat()
+                    preserved_configs[cfg_file.name] = (
+                        cfg_file.read_bytes(),
+                        st.st_mode,
+                        st.st_atime,
+                        st.st_mtime,
+                    )
 
+        # Install extension (dest_dir computed above during self-install guard).
+        # Preserved configs are the ONLY surviving copy once the old extension dir
+        # is removed, so guard the whole destructive reinstall against loss at any
+        # step (a malformed source ignore file, the rmtree, copytree, or a
+        # restore):
+        #   1. Validate the source ignore file BEFORE deleting anything — it can
+        #      raise on a malformed file, and nothing is destroyed yet if it does.
+        #   2. Stage a durable on-disk backup of the configs.
+        #   3. Run the destructive rmtree + copytree + restore inside one block;
+        #      on ANY failure, roll the configs back into a clean dest_dir and
+        #      re-raise. The backup is removed only once the configs are provably
+        #      in place (normal success OR a completed rollback) — otherwise it is
+        #      left on disk so nothing is ever lost, even on a rollback failure.
+        # Each config write is atomic (temp file + os.replace) so it never follows
+        # a symlink that may occupy the path (see _atomic_restore_config).
         ignore_fn = self._load_extensionignore(source_dir)
-        shutil.copytree(source_dir, dest_dir, ignore=ignore_fn)
+
+        backup_dir: Path | None = None
+        if preserved_configs:
+            backup_dir = Path(
+                tempfile.mkdtemp(prefix=f".{dest_dir.name}.cfgbak-", dir=dest_dir.parent)
+            )
+            for name, (data, mode, atime, mtime) in preserved_configs.items():
+                self._atomic_restore_config(backup_dir / name, data, mode, atime, mtime)
+
+        restored_ok = not preserved_configs
+        try:
+            if dest_dir.exists():
+                shutil.rmtree(dest_dir)
+            shutil.copytree(source_dir, dest_dir, ignore=ignore_fn)
+            # Restore every preserved config. A failure here is NOT swallowed: a
+            # config we promised to preserve but could not restore must fail the
+            # reinstall, not report success with the config silently missing.
+            for name, (data, mode, atime, mtime) in preserved_configs.items():
+                self._atomic_restore_config(dest_dir / name, data, mode, atime, mtime)
+            restored_ok = True
+        except Exception:
+            # Roll the configs back into a clean dest_dir from the durable backup
+            # so a failed reinstall leaves them intact, then re-raise the original.
+            if backup_dir is not None:
+                try:
+                    if dest_dir.exists():
+                        shutil.rmtree(dest_dir, ignore_errors=True)
+                    shutil.copytree(backup_dir, dest_dir, dirs_exist_ok=True)
+                    restored_ok = True
+                except OSError:
+                    restored_ok = False  # keep the backup below for recovery
+            raise
+        finally:
+            # Drop the backup only when the configs are provably in dest_dir.
+            if backup_dir is not None and restored_ok and backup_dir.exists():
+                shutil.rmtree(backup_dir, ignore_errors=True)
 
         # Register commands with AI agents
         registered_commands = {}

tests/test_extensions.py

@@ -1557,6 +1557,235 @@ def test_config_backup_on_remove(self, extension_dir, project_dir):
         assert backup_file.exists()
         assert backup_file.read_text() == "test: config"
 
+    def test_reinstall_preserves_keep_config_leftover(self, extension_dir, project_dir):
+        """A reinstall after `remove(keep_config=True)` must not wipe the
+        preserved config. remove(keep_config=True) leaves the *-config.yml in
+        place and drops the registry entry; the subsequent reinstall is not a
+        --force removal, so the unconditional rmtree used to destroy it."""
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+        # Restrictive perms — a config may hold secrets; the restore must not
+        # widen them (POSIX only; Windows chmod ignores group/other bits).
+        config_file.chmod(0o600)
+        # Distinctive mtime to prove timestamps survive too (mirrors copy2).
+        old_mtime = 1_600_000_000
+        os.utime(config_file, (old_mtime, old_mtime))
+
+        # keep-config removal: registry entry dropped, config left in place.
+        assert manager.remove("test-ext", keep_config=True) is True
+        assert config_file.exists()
+
+        # Reinstall (no --force; registry now reports not-installed).
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        # The user's config must survive the reinstall (was silently wiped).
+        assert config_file.exists()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+        # ...and keep its original mtime (copy2-style timestamp preservation).
+        assert config_file.stat().st_mtime == old_mtime
+        # ...and its original (restrictive) mode, not default perms.
+        if platform.system() != "Windows":
+            assert config_file.stat().st_mode & 0o777 == 0o600
+
+    def test_reinstall_config_restore_does_not_follow_symlink(
+        self, extension_dir, project_dir, monkeypatch
+    ):
+        """The keep-config restore must never write *through* a symlink. If the
+        config path is a symlink at restore time (e.g. swapped in between the
+        copytree and the restore), the write must not follow it and clobber the
+        external target — it must replace the link with a fresh regular file."""
+        # Probe symlink capability (Windows usually needs privilege).
+        probe_target = project_dir / "_probe_target"
+        probe_link = project_dir / "_probe_link"
+        try:
+            probe_target.write_text("x")
+            probe_link.symlink_to(probe_target)
+        except (OSError, NotImplementedError):
+            pytest.skip("symlinks not supported in this environment")
+        finally:
+            probe_link.unlink(missing_ok=True)
+            probe_target.unlink(missing_ok=True)
+
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+
+        # keep-config removal: registry entry dropped, config left in place.
+        assert manager.remove("test-ext", keep_config=True) is True
+
+        # A victim file outside the extension dir that the symlink points at.
+        victim = project_dir / "victim.txt"
+        victim.write_text("DO-NOT-OVERWRITE")
+
+        # Simulate the config path being a symlink to the victim right after the
+        # fresh copytree (the TOCTOU / symlinks=True case the guard defends).
+        real_copytree = shutil.copytree
+
+        def copytree_then_symlink(src, dst, *args, **kwargs):
+            result = real_copytree(src, dst, *args, **kwargs)
+            link = Path(dst) / "test-ext-config.yml"
+            if link.exists() or link.is_symlink():
+                link.unlink()
+            link.symlink_to(victim)
+            return result
+
+        monkeypatch.setattr(shutil, "copytree", copytree_then_symlink)
+
+        # Reinstall: restore must drop the symlink, not follow it.
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        # The victim outside the extension dir is untouched.
+        assert victim.read_text() == "DO-NOT-OVERWRITE"
+        # The config path is now a regular file (not a symlink) with the
+        # preserved content written in place.
+        assert not config_file.is_symlink()
+        assert config_file.is_file()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+
+    def test_reinstall_config_restore_replaces_directory_at_path(
+        self, extension_dir, project_dir, monkeypatch
+    ):
+        """Cross-platform guard check: if a directory occupies the config path at
+        restore time, write_bytes() would raise (uncaught, crashing the whole
+        reinstall). The guard must clear the directory and write a fresh regular
+        file. Exercises the same non-regular-file guard as the symlink case, but
+        needs no symlink privilege so it runs everywhere (incl. Windows)."""
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+        assert manager.remove("test-ext", keep_config=True) is True
+
+        # Simulate a directory occupying the config path right after copytree.
+        real_copytree = shutil.copytree
+
+        def copytree_then_dir(src, dst, *args, **kwargs):
+            result = real_copytree(src, dst, *args, **kwargs)
+            clash = Path(dst) / "test-ext-config.yml"
+            clash.mkdir()
+            (clash / "sentinel").write_text("stray")
+            return result
+
+        monkeypatch.setattr(shutil, "copytree", copytree_then_dir)
+
+        # Must not raise, and must restore a regular file with preserved content.
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        assert config_file.is_file()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+
+    def test_reinstall_config_survives_copytree_failure(
+        self, extension_dir, project_dir, monkeypatch
+    ):
+        """The preserved config is captured in memory and the old dir is removed
+        before copytree runs. If copytree fails, the config must be written back
+        (not lost) so a failed reinstall leaves the user's config intact."""
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+        config_file.chmod(0o600)
+        assert manager.remove("test-ext", keep_config=True) is True
+
+        # Fail the source→dest copytree (after the old dir is removed) but let the
+        # rollback copytree (backup→dest) succeed, so recovery can be observed.
+        real_copytree = shutil.copytree
+        calls = {"n": 0}
+
+        def boom(src, dst, *args, **kwargs):
+            calls["n"] += 1
+            if calls["n"] == 1:
+                raise OSError("simulated disk-full during copytree")
+            return real_copytree(src, dst, *args, **kwargs)
+
+        monkeypatch.setattr(shutil, "copytree", boom)
+
+        with pytest.raises(OSError):
+            manager.install_from_directory(
+                extension_dir, "0.1.0", register_commands=False
+            )
+
+        # Config recovered despite the failure (was permanently lost before).
+        assert config_file.is_file()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+        if platform.system() != "Windows":
+            assert config_file.stat().st_mode & 0o777 == 0o600
+
+    def test_reinstall_validates_ignore_before_deleting_config(
+        self, extension_dir, project_dir, monkeypatch
+    ):
+        """A malformed source ignore file must be detected BEFORE the destructive
+        rmtree, so a reinstall that fails on it leaves the preserved config
+        untouched (it used to be loaded only after rmtree had run)."""
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+        assert manager.remove("test-ext", keep_config=True) is True
+
+        def bad_ignore(src):
+            raise ValueError("malformed .extensionignore")
+
+        monkeypatch.setattr(manager, "_load_extensionignore", bad_ignore)
+
+        with pytest.raises(ValueError):
+            manager.install_from_directory(
+                extension_dir, "0.1.0", register_commands=False
+            )
+
+        # Nothing was deleted — the leftover config is still in place.
+        assert config_file.is_file()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+
+    def test_reinstall_failed_config_restore_is_not_silent(
+        self, extension_dir, project_dir, monkeypatch
+    ):
+        """A restore failure after copytree must NOT be silently skipped and
+        reported as a successful install — it fails loudly, and the config is
+        rolled back from the durable backup rather than left missing."""
+        manager = ExtensionManager(project_dir)
+        manager.install_from_directory(extension_dir, "0.1.0", register_commands=False)
+
+        ext_dir = project_dir / ".specify" / "extensions" / "test-ext"
+        config_file = ext_dir / "test-ext-config.yml"
+        config_file.write_text("api_key: MY-CUSTOMIZED-VALUE")
+        assert manager.remove("test-ext", keep_config=True) is True
+
+        # Fail the live restore into the extension dir, but let the durable-backup
+        # staging (a ".cfgbak-" temp dir) succeed so rollback can recover.
+        real_restore = ExtensionManager._atomic_restore_config
+
+        def flaky_restore(dest, data, mode, atime, mtime):
+            if "cfgbak" not in dest.parent.name:
+                raise OSError("simulated restore failure")
+            return real_restore(dest, data, mode, atime, mtime)
+
+        monkeypatch.setattr(
+            ExtensionManager, "_atomic_restore_config", staticmethod(flaky_restore)
+        )
+
+        with pytest.raises(OSError):
+            manager.install_from_directory(
+                extension_dir, "0.1.0", register_commands=False
+            )
+
+        # Not a silent success: the config was rolled back, not lost.
+        assert config_file.is_file()
+        assert "MY-CUSTOMIZED-VALUE" in config_file.read_text()
+
 
 # ===== CommandRegistrar Tests =====