diff --git a/src/specify_cli/extensions/__init__.py b/src/specify_cli/extensions/__init__.py index e00b9c247a..3d75e8c5dc 100644 --- a/src/specify_cli/extensions/__init__.py +++ b/src/specify_cli/extensions/__init__.py @@ -14,6 +14,7 @@ import os import re import shutil +import stat import tempfile import zipfile from dataclasses import dataclass @@ -101,6 +102,34 @@ def _load_core_command_names() -> frozenset[str]: CORE_COMMAND_NAMES = _load_core_command_names() +def _fsync_file(path: Path) -> None: + """Best-effort fsync for a regular file.""" + try: + with open(path, "rb") as handle: + os.fsync(handle.fileno()) + except (AttributeError, OSError, NotImplementedError): + pass + + +def _fsync_directory(path: Path) -> None: + """Best-effort fsync for a directory path.""" + if not path.exists(): + return + try: + dir_fd = os.open(str(path), os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)) + except (AttributeError, OSError, NotImplementedError): + try: + dir_fd = os.open(str(path), os.O_RDONLY) + except (AttributeError, OSError, NotImplementedError): + return + try: + os.fsync(dir_fd) + except (AttributeError, OSError, NotImplementedError): + pass + finally: + os.close(dir_fd) + + class ExtensionError(Exception): """Base exception for extension-related errors.""" @@ -1402,12 +1431,188 @@ def install_from_directory( backup_config_dir.unlink() did_remove = self.remove(manifest.id) + # Rescue any config files left behind by a prior `remove --keep-config`. + # When an extension is removed with --keep-config, it is no longer in + # the registry but its config files remain in dest_dir. A subsequent + # plain (non-force) install would delete that directory unconditionally, + # silently discarding the preserved config. We read those files into + # memory and also write a durable staging copy outside dest_dir so + # that a partial rmtree, failed copytree, or partial restore cannot + # permanently discard the user's original bytes on a retry. The + # staging dir is removed only after every config has been successfully + # restored. + stranded_configs: dict[str, tuple[bytes, int]] = {} + rescue_staging_dir = self.extensions_dir / f".rescue-staging-{manifest.id}" + # A staging directory is trusted only when this completion marker is + # present. The marker is written after every staged file is complete + # and removed before the non-atomic cleanup, so a crash mid-staging or + # mid-cleanup can never leave a partial directory that a retry mistakes + # for a complete durable backup. + rescue_complete_marker = rescue_staging_dir / ".rescue-complete" + staging_is_complete = ( + rescue_staging_dir.is_dir() + and not rescue_staging_dir.is_symlink() + and rescue_complete_marker.is_file() + and not rescue_complete_marker.is_symlink() + ) + + if staging_is_complete and not self.registry.is_installed(manifest.id): + # A previous install attempt staged the configs but never + # completed cleanly. Reload from the durable backup so the + # original bytes are used on retry rather than whatever + # mixture of packaged defaults and partial restores remains + # on disk. Only load non-symlinked files whose names match + # the two recognised config suffixes so a tampered staging + # directory cannot inject arbitrary files. + for staged_file in sorted(rescue_staging_dir.iterdir()): + if ( + staged_file.is_file() + and not staged_file.is_symlink() + and staged_file.name.endswith(("-config.yml", "-config.local.yml")) + ): + stranded_configs[staged_file.name] = ( + staged_file.read_bytes(), + staged_file.stat().st_mode, + ) + elif dest_dir.exists() and not self.registry.is_installed(manifest.id): + for cfg_file in ( + list(dest_dir.glob("*-config.yml")) + + list(dest_dir.glob("*-config.local.yml")) + ): + if cfg_file.is_file() and not cfg_file.is_symlink(): + stranded_configs[cfg_file.name] = ( + cfg_file.read_bytes(), + cfg_file.stat().st_mode, + ) + + if stranded_configs and not staging_is_complete: + # Write a durable backup outside dest_dir before any + # destructive operation so the original bytes survive a + # crash or partial failure at any later step. The staging + # dir is cleaned up only after every restore succeeds. + # + # Any pre-existing staging dir here lacks the completion marker + # (staging_is_complete is False), so it is a stale partial from an + # interrupted attempt — remove it first for a clean write. + if rescue_staging_dir.is_symlink(): + rescue_staging_dir.unlink() + elif rescue_staging_dir.is_dir(): + shutil.rmtree(rescue_staging_dir) + elif rescue_staging_dir.exists(): + rescue_staging_dir.unlink() + try: + rescue_staging_dir.mkdir(parents=True, exist_ok=True) + for filename, (content, mode) in stranded_configs.items(): + staged = rescue_staging_dir / filename + # Create the staging file with mode 0600 before writing so + # the preserved bytes are never transiently readable by other + # local users, even on a umask that would produce 0644. + fd = os.open(str(staged), os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) + try: + # os.write() may write fewer bytes than requested, so + # loop until the whole buffer is on disk — a truncated + # "durable" backup would be trusted over the intact + # config on a retry and cause silent data loss. + view = memoryview(content) + written = 0 + while written < len(view): + written += os.write(fd, view[written:]) + try: + os.fchmod(fd, stat.S_IMODE(mode)) + except (AttributeError, NotImplementedError, OSError): + try: + staged.chmod(stat.S_IMODE(mode)) + except (NotImplementedError, OSError): + pass # Best-effort; chmod may not be supported on all platforms. + try: + os.fsync(fd) + except (AttributeError, OSError, NotImplementedError): + pass + finally: + os.close(fd) + # Write the completion marker only after every staged file is + # fully written so a retry trusts staging only when it is whole. + marker_fd = os.open( + str(rescue_complete_marker), + os.O_WRONLY | os.O_CREAT | os.O_EXCL, + 0o600, + ) + try: + os.fsync(marker_fd) + except (AttributeError, OSError, NotImplementedError): + pass + finally: + os.close(marker_fd) + _fsync_directory(rescue_staging_dir) + _fsync_directory(rescue_staging_dir.parent) + except BaseException: + # Durable staging failed (or was interrupted). Continuing with + # only the in-memory copy would reintroduce the permanent-loss + # path this staging exists to close: the rmtree below could + # delete the originals and a later restore failure would leave + # no on-disk copy. dest_dir is still untouched here, so clean + # up the partial staging dir and abort the install instead of + # proceeding destructively. + shutil.rmtree(rescue_staging_dir, ignore_errors=True) + raise + # Install extension (dest_dir computed above during self-install guard) if dest_dir.exists(): shutil.rmtree(dest_dir) ignore_fn = self._load_extensionignore(source_dir) - shutil.copytree(source_dir, dest_dir, ignore=ignore_fn) + + def _restore_stranded_config_file( + target: Path, content: bytes, preserved_mode: int + ) -> None: + tmp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=target.parent, + prefix=f".{target.name}.", + delete=False, + ) as tmp: + tmp_path = Path(tmp.name) + tmp.write(content) + try: + tmp_path.chmod(stat.S_IMODE(preserved_mode)) + except (NotImplementedError, OSError): + pass # Best-effort; chmod may not be supported on all platforms. + os.replace(tmp_path, target) + except BaseException: + if tmp_path is not None and tmp_path.exists(): + tmp_path.unlink() + raise + + try: + shutil.copytree(source_dir, dest_dir, ignore=ignore_fn) + except BaseException: + # copytree failed — dest_dir may be absent or only partially + # created. Write the rescued configs back now so they are not + # permanently lost even though the install did not complete. + if stranded_configs: + dest_dir.mkdir(parents=True, exist_ok=True) + for filename, (content, mode) in stranded_configs.items(): + target = dest_dir / filename + _restore_stranded_config_file(target, content, mode) + raise + + # Restore stranded configs rescued before the rmtree above. + for filename, (content, mode) in stranded_configs.items(): + target = dest_dir / filename + _restore_stranded_config_file(target, content, mode) + + # Every stranded config has been restored successfully — the + # durable staging backup is no longer needed. Raise on failure so + # the install is not reported as successful while a stale backup + # that could be misread on the next retry remains on disk. + if rescue_staging_dir.is_dir() and not rescue_staging_dir.is_symlink(): + # Remove the completion marker before the non-atomic rmtree so a + # crash mid-cleanup cannot leave a staging dir that a retry would + # wrongly trust as a complete durable backup. + rescue_complete_marker.unlink(missing_ok=True) + shutil.rmtree(rescue_staging_dir) # Register commands with AI agents registered_commands = {} diff --git a/tests/test_extensions.py b/tests/test_extensions.py index d2f4af8264..9fbc20340d 100644 --- a/tests/test_extensions.py +++ b/tests/test_extensions.py @@ -20,6 +20,7 @@ from pathlib import Path from datetime import datetime, timezone from unittest.mock import MagicMock +from specify_cli import extensions as _ext_module from tests.conftest import strip_ansi from specify_cli.extensions import ( @@ -1226,6 +1227,202 @@ def test_install_force_config_preserved(self, extension_dir, project_dir): assert new_config.exists() assert new_config.read_text() == "test: config" + def test_reinstall_after_keep_config_preserves_config( + self, extension_dir, project_dir + ): + """Reinstalling after `remove --keep-config` must not overwrite preserved config.""" + manager = ExtensionManager(project_dir) + + # Add a packaged default config so the reinstall has a file to overwrite. + # Without the fix, the packaged default would silently win on reinstall. + packaged_config = extension_dir / "test-ext-config.yml" + packaged_config.write_text("model: default-model\nmax_iterations: 1\n") + + # Install once (packaged default is copied into the installed directory) + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + # Overwrite the installed config with user-customized values + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + config_file = ext_dir / "test-ext-config.yml" + config_file.write_text("model: custom-model\nmax_iterations: 99\n") + + # Remove while preserving config + manager.remove("test-ext", keep_config=True) + assert not manager.registry.is_installed("test-ext") + assert config_file.exists() + assert "custom-model" in config_file.read_text() + + # Plain reinstall (no --force) — packaged default is still present in + # extension_dir, so a naive implementation would overwrite the custom values. + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + # Preserved config must survive the reinstall and beat the packaged default + assert config_file.exists() + assert "custom-model" in config_file.read_text() + assert "99" in config_file.read_text() + assert "default-model" not in config_file.read_text() + + def test_reinstall_after_keep_config_preserves_local_config( + self, extension_dir, project_dir + ): + """Local config override files (*-config.local.yml) are also rescued on reinstall.""" + manager = ExtensionManager(project_dir) + + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + ext_dir = project_dir / ".specify" / "extensions" / "test-ext" + local_cfg = ext_dir / "test-ext-config.local.yml" + local_cfg.write_text("local_override: true\n") + + manager.remove("test-ext", keep_config=True) + assert local_cfg.exists() + + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + assert local_cfg.exists() + assert "local_override: true" in local_cfg.read_text() + + def test_copytree_failure_restores_stranded_config( + self, extension_dir, project_dir, monkeypatch + ): + """A copytree failure must not permanently lose a preserved config. + + When copytree raises after the existing directory has been removed, the + rollback path must write the rescued bytes back to dest_dir and restore + the original file mode, while leaving the extension unregistered. + """ + import stat + + manager = ExtensionManager(project_dir) + + # Add a packaged default config so copytree would overwrite it on success. + packaged_config = extension_dir / "test-ext-config.yml" + packaged_config.write_text("model: default-model\n") + + # Install once so the extension is on disk. + 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("model: custom-model\nmax_iterations: 99\n") + + # Set a known, non-default mode so we can assert it survives the rollback. + if platform.system() != "Windows": + config_file.chmod(0o640) + original_bytes = config_file.read_bytes() + original_mode = config_file.stat().st_mode + + # Remove while preserving the config — it is now a stranded file. + manager.remove("test-ext", keep_config=True) + assert not manager.registry.is_installed("test-ext") + assert config_file.exists() + + # Make copytree create a partial destination then raise so the rollback + # path is exercised. + + def failing_copytree(src, dst, **kwargs): + Path(dst).mkdir(parents=True, exist_ok=True) + (Path(dst) / "_partial.txt").write_text("partial") + raise OSError("simulated disk full") + + monkeypatch.setattr(_ext_module.shutil, "copytree", failing_copytree) + + with pytest.raises(OSError, match="simulated disk full"): + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + # The preserved config must have been written back by the rollback path. + assert config_file.exists(), "rollback must recreate the config file" + assert config_file.read_bytes() == original_bytes + + # On POSIX, the original file mode must be faithfully restored. + if platform.system() != "Windows": + restored_mode = config_file.stat().st_mode + assert stat.S_IMODE(restored_mode) == stat.S_IMODE(original_mode) + + # The extension must remain unregistered after the failed install. + assert not manager.registry.is_installed("test-ext") + + def test_retry_after_staging_backup_restores_stranded_config( + self, extension_dir, project_dir, monkeypatch + ): + """A retry after an interrupted install should restore the rescued config from staging.""" + import stat + + manager = ExtensionManager(project_dir) + + packaged_config = extension_dir / "test-ext-config.yml" + packaged_config.write_text("model: default-model\n") + + 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("model: custom-model\nmax_iterations: 99\n") + + if platform.system() != "Windows": + config_file.chmod(0o640) + original_bytes = config_file.read_bytes() + original_mode = config_file.stat().st_mode + + manager.remove("test-ext", keep_config=True) + assert not manager.registry.is_installed("test-ext") + assert config_file.exists() + + staging_dir = manager.extensions_dir / ".rescue-staging-test-ext" + assert not staging_dir.exists() + + original_copytree = shutil.copytree + copytree_calls = 0 + + def flaky_copytree(*args, **kwargs): + nonlocal copytree_calls + copytree_calls += 1 + if copytree_calls == 1: + dst = args[1] + Path(dst).mkdir(parents=True, exist_ok=True) + (Path(dst) / "_partial.txt").write_text("partial") + raise OSError("simulated disk full") + return original_copytree(*args, **kwargs) + + monkeypatch.setattr(_ext_module.shutil, "copytree", flaky_copytree) + + with pytest.raises(OSError, match="simulated disk full"): + manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + assert staging_dir.exists() + assert (staging_dir / ".rescue-complete").exists() + assert (staging_dir / "test-ext-config.yml").exists() + + manifest = manager.install_from_directory( + extension_dir, "0.1.0", register_commands=False + ) + + assert manifest.id == "test-ext" + assert manager.registry.is_installed("test-ext") + assert config_file.read_bytes() == original_bytes + assert not (ext_dir / "_partial.txt").exists() + assert not staging_dir.exists() + + if platform.system() != "Windows": + restored_mode = config_file.stat().st_mode + assert stat.S_IMODE(restored_mode) == stat.S_IMODE(original_mode) + def test_install_force_without_existing(self, extension_dir, project_dir): """Test force-install when extension is NOT already installed (works normally).""" manager = ExtensionManager(project_dir)