Address review: collision-resistant module names, extra_files support, remove orphan dir

Author: Copilot

src/specify_cli/__init__.py

@@ -6097,12 +6097,43 @@ def _safe_fetch(url: str) -> bytes:
             )
             raise typer.Exit(1)
 
-        # Write validated files to the staged temp location, then atomically rename
-        # into place. Both paths are under steps_base_dir (same filesystem), so
-        # os.rename() is atomic on POSIX and won't leave a partially-written
-        # directory at step_dir on failure.
+        # Write the two required files.
         (tmp_path / "step.yml").write_bytes(step_yml_content)
         (tmp_path / "__init__.py").write_bytes(init_py_content)
+
+        # Optionally download additional package files declared in the catalog entry
+        # (e.g. helper modules). Each entry in ``extra_files`` is a mapping of
+        # relative-path → URL. step.yml and __init__.py are ignored here (already
+        # written). Paths are validated to stay within the step package directory to
+        # prevent path-traversal attacks.
+        extra_files = info.get("extra_files") or {}
+        if isinstance(extra_files, dict):
+            for rel_path, file_url in extra_files.items():
+                if rel_path in ("step.yml", "__init__.py"):
+                    continue  # already written above
+                dest = (tmp_path / rel_path).resolve()
+                try:
+                    dest.relative_to(tmp_path)
+                except ValueError:
+                    console.print(
+                        f"[red]Error:[/red] extra_files path '{rel_path}' is outside "
+                        "the step package directory"
+                    )
+                    raise typer.Exit(1)
+                try:
+                    file_content = _safe_fetch(file_url)
+                except Exception as exc:
+                    console.print(
+                        f"[red]Error:[/red] Failed to download extra file '{rel_path}': {exc}"
+                    )
+                    raise typer.Exit(1)
+                dest.parent.mkdir(parents=True, exist_ok=True)
+                dest.write_bytes(file_content)
+
+        # Atomically rename the staging directory to the final location.
+        # Both paths are under steps_base_dir (same filesystem), so os.rename()
+        # is atomic on POSIX and won't leave a partially-written directory at
+        # step_dir on failure.
         os.rename(tmp_path, step_dir)
     finally:
         # Clean up if the rename hasn't moved tmp_path yet (i.e. on any failure).
@@ -6145,9 +6176,7 @@ def workflow_step_remove(
         raise typer.Exit(1)
 
     registry = StepRegistry(project_root)
-    if not registry.is_installed(step_id):
-        console.print(f"[red]Error:[/red] Step type '{step_id}' is not installed")
-        raise typer.Exit(1)
+    in_registry = registry.is_installed(step_id)
 
     steps_base_dir = (project_root / ".specify" / "workflows" / "steps").resolve()
     step_dir = (steps_base_dir / step_id).resolve()
@@ -6156,11 +6185,28 @@ def workflow_step_remove(
     except ValueError:
         console.print(f"[red]Error:[/red] Invalid step id '{step_id}'")
         raise typer.Exit(1)
-    if step_dir.exists():
+
+    dir_exists = step_dir.exists()
+
+    if not in_registry and not dir_exists:
+        console.print(f"[red]Error:[/red] Step type '{step_id}' is not installed")
+        raise typer.Exit(1)
+
+    if not in_registry and dir_exists:
+        # The registry was likely reset due to corruption.  Warn the user that the
+        # directory is being removed even though there is no registry entry, so
+        # the orphaned package can be cleaned up and a fresh install attempted.
+        console.print(
+            f"[yellow]Warning:[/yellow] '{step_id}' has no registry entry "
+            "(registry may have been reset). Removing the orphaned directory."
+        )
+
+    if dir_exists:
         import shutil
         shutil.rmtree(step_dir)
 
-    registry.remove(step_id)
+    if in_registry:
+        registry.remove(step_id)
     console.print(f"[green]✓[/green] Step type '{step_id}' uninstalled")
 
 

src/specify_cli/workflows/__init__.py

@@ -80,6 +80,7 @@ def load_custom_steps(project_root: Path) -> list[str]:
     Returns a list of type_keys that were successfully loaded.
     Silently skips packages that fail to import or validate.
     """
+    import hashlib as _hashlib
     import importlib.util as _importlib_util
     import re as _re
     import sys as _sys
@@ -111,9 +112,14 @@ def load_custom_steps(project_root: Path) -> list[str]:
                 continue
 
             # Sanitize type_key so the synthetic module name is a valid identifier
-            # (e.g. "test-custom" → "_speckit_custom_step_test_custom").
+            # (e.g. "test-custom" → "_speckit_custom_step_test_custom_<hash>").
+            # The 8-char SHA-256 hash of the original type_key makes the name
+            # collision-resistant when different type_keys produce the same
+            # sanitized form (e.g. "a-b" and "a_b" both sanitize to "a_b" but
+            # have different hashes).
             safe_key = _re.sub(r"[^A-Za-z0-9_]", "_", type_key)
-            module_name = f"_speckit_custom_step_{safe_key}"
+            key_hash = _hashlib.sha256(type_key.encode()).hexdigest()[:8]
+            module_name = f"_speckit_custom_step_{safe_key}_{key_hash}"
 
             # Treat the step directory as a proper package so that relative
             # imports inside the step (e.g. ``from .helpers import …``) work.

tests/test_workflows.py

@@ -2256,6 +2256,7 @@ def test_skip_broken_init_py(self, project_dir):
 
     def test_module_name_sanitized_for_hyphenated_type_key(self, project_dir):
         """type_key values with hyphens produce valid Python module identifiers."""
+        import hashlib
         import sys
         from specify_cli.workflows import load_custom_steps, STEP_REGISTRY
 
@@ -2281,10 +2282,14 @@ def execute(self, config, context):
         assert "my-hyphen-step" in loaded
         assert "my-hyphen-step" in STEP_REGISTRY
         # Synthetic module name must be a valid identifier (hyphens → underscores)
-        assert "_speckit_custom_step_my_hyphen_step" in sys.modules
+        # and include a collision-resistant hash suffix.
+        key_hash = hashlib.sha256(b"my-hyphen-step").hexdigest()[:8]
+        module_name = f"_speckit_custom_step_my_hyphen_step_{key_hash}"
+        assert module_name in sys.modules
 
     def test_package_relative_import(self, project_dir):
         """Steps can use relative imports to access sibling modules."""
+        import hashlib
         import sys
         from specify_cli.workflows import load_custom_steps, STEP_REGISTRY
 
@@ -2314,7 +2319,94 @@ def execute(self, config, context):
         loaded = load_custom_steps(project_dir)
         assert "pkg-step" in loaded
         assert "pkg-step" in STEP_REGISTRY
-        # Verify the relative import actually resolved
-        module_name = "_speckit_custom_step_pkg_step"
+        # Verify the relative import actually resolved; module name includes hash suffix.
+        key_hash = hashlib.sha256(b"pkg-step").hexdigest()[:8]
+        module_name = f"_speckit_custom_step_pkg_step_{key_hash}"
         assert module_name in sys.modules
         assert sys.modules[module_name].PkgStep.helper == "hello"
+
+    def test_module_name_collision_resistance(self, project_dir):
+        """'a-b' and 'a_b' produce different module names despite the same sanitized form."""
+        import hashlib
+
+        # Simulate the module name generation for two type_keys that sanitize the same way
+        def make_module_name(type_key: str) -> str:
+            import re
+            safe_key = re.sub(r"[^A-Za-z0-9_]", "_", type_key)
+            key_hash = hashlib.sha256(type_key.encode()).hexdigest()[:8]
+            return f"_speckit_custom_step_{safe_key}_{key_hash}"
+
+        name_a = make_module_name("a-b")
+        name_b = make_module_name("a_b")
+        assert name_a != name_b, "Module names for 'a-b' and 'a_b' must differ"
+
+
+# ===== CLI Step Remove Tests =====
+
+class TestWorkflowStepRemoveCLI:
+    """Test the 'specify workflow step remove' CLI command edge cases."""
+
+    def test_remove_orphaned_directory(self, project_dir, monkeypatch):
+        """step remove works when directory exists but registry entry is missing.
+
+        This covers the case where the registry was reset due to corruption.
+        """
+        from typer.testing import CliRunner
+        from specify_cli import app
+
+        monkeypatch.chdir(project_dir)
+
+        # Create an orphaned step directory (no registry entry)
+        step_dir = project_dir / ".specify" / "workflows" / "steps" / "orphan-step"
+        step_dir.mkdir(parents=True)
+        (step_dir / "step.yml").write_text(
+            "step:\n  type_key: orphan-step\n", encoding="utf-8"
+        )
+        (step_dir / "__init__.py").write_text("", encoding="utf-8")
+
+        runner = CliRunner()
+        result = runner.invoke(app, ["workflow", "step", "remove", "orphan-step"])
+
+        assert result.exit_code == 0, result.output
+        assert not step_dir.exists()
+        # Warning should be printed about missing registry entry
+        assert "Warning" in result.output or "warning" in result.output.lower()
+
+    def test_remove_not_installed(self, project_dir, monkeypatch):
+        """step remove fails cleanly when neither directory nor registry entry exist."""
+        from typer.testing import CliRunner
+        from specify_cli import app
+
+        monkeypatch.chdir(project_dir)
+
+        runner = CliRunner()
+        result = runner.invoke(app, ["workflow", "step", "remove", "ghost-step"])
+
+        assert result.exit_code != 0
+        assert "not installed" in result.output
+
+    def test_remove_registered_step(self, project_dir, monkeypatch):
+        """step remove works normally when both directory and registry entry exist."""
+        from typer.testing import CliRunner
+        from specify_cli import app
+        from specify_cli.workflows.catalog import StepRegistry
+
+        monkeypatch.chdir(project_dir)
+
+        # Set up a registered step with a directory
+        registry = StepRegistry(project_dir)
+        registry.add("my-step", {"name": "My Step", "type_key": "my-step", "version": "1.0.0"})
+        step_dir = project_dir / ".specify" / "workflows" / "steps" / "my-step"
+        step_dir.mkdir(parents=True)
+        (step_dir / "step.yml").write_text(
+            "step:\n  type_key: my-step\n", encoding="utf-8"
+        )
+        (step_dir / "__init__.py").write_text("", encoding="utf-8")
+
+        runner = CliRunner()
+        result = runner.invoke(app, ["workflow", "step", "remove", "my-step"])
+
+        assert result.exit_code == 0, result.output
+        assert not step_dir.exists()
+        registry2 = StepRegistry(project_dir)
+        assert not registry2.is_installed("my-step")