fix: model list and remove now find files in subdirectories

Author: bigcat88

comfy_cli/command/models/models.py

@@ -358,14 +358,19 @@ def remove(
         typer.echo("No models found to remove.")
         return
 
+    model_dir_resolved = model_dir.resolve()
+
     to_delete = []
     # Scenario #1: User provided model names to delete
     if model_names:
         # Validate and filter models to delete based on provided names
         missing_models = []
         for name in model_names:
-            model_path = model_dir / name
-            if model_path.exists():
+            model_path = (model_dir / name).resolve()
+            if not model_path.is_relative_to(model_dir_resolved):
+                typer.echo(f"Invalid model path: {name}")
+                continue
+            if model_path.is_file():
                 to_delete.append(model_path)
             else:
                 missing_models.append(name)
@@ -377,7 +382,8 @@ def remove(
 
     # Scenario #2: User did not provide model names, prompt for selection
     else:
-        selections = ui.prompt_multi_select("Select models to delete:", [model.name for model in available_models])
+        rel_names = [str(model.relative_to(model_dir)) for model in available_models]
+        selections = ui.prompt_multi_select("Select models to delete:", rel_names)
         if not selections:
             typer.echo("No models selected for deletion.")
             return
@@ -394,9 +400,11 @@ def remove(
         typer.echo("Deletion canceled.")
 
 
-def list_models(path: pathlib.Path) -> list:
-    """List all models in the specified directory."""
-    return [file for file in path.iterdir() if file.is_file()]
+def list_models(path: pathlib.Path) -> list[pathlib.Path]:
+    """List all model files recursively in the specified directory."""
+    if not path.is_dir():
+        return []
+    return sorted(f for f in path.rglob("*") if f.is_file())
 
 
 @app.command("list")
@@ -418,6 +426,10 @@ def list_command(
         return
 
     # Prepare data for table display
-    data = [(model.name, f"{model.stat().st_size // 1024} KB") for model in models]
-    column_names = ["Model Name", "Size"]
+    data = []
+    for model in models:
+        rel = model.relative_to(model_dir)
+        model_type = str(rel.parent) if len(rel.parts) > 1 else ""
+        data.append((model.name, model_type, f"{model.stat().st_size // 1024} KB"))
+    column_names = ["Model Name", "Type", "Size"]
     ui.display_table(data, column_names)

tests/comfy_cli/command/models/test_models.py

@@ -1,4 +1,152 @@
-from comfy_cli.command.models.models import check_civitai_url, check_huggingface_url
+import pathlib
+from unittest.mock import patch
+
+import typer.testing
+
+from comfy_cli.command.models.models import app, check_civitai_url, check_huggingface_url, list_models
+
+
+def _make_model_tree(tmp_path: pathlib.Path) -> pathlib.Path:
+    """Create a realistic model directory tree and return its root."""
+    model_dir = tmp_path / "models"
+    # Root-level file
+    (model_dir / "root_model.safetensors").parent.mkdir(parents=True, exist_ok=True)
+    (model_dir / "root_model.safetensors").write_bytes(b"x" * 100)
+    # One level deep (checkpoints/)
+    (model_dir / "checkpoints").mkdir()
+    (model_dir / "checkpoints" / "sd15.safetensors").write_bytes(b"x" * 200)
+    # Two levels deep (loras/SD1.5/)
+    (model_dir / "loras" / "SD1.5").mkdir(parents=True)
+    (model_dir / "loras" / "SD1.5" / "detail.safetensors").write_bytes(b"x" * 300)
+    # An empty subdirectory (should be ignored)
+    (model_dir / "empty_dir").mkdir()
+    return model_dir
+
+
+# ---------------------------------------------------------------------------
+# list_models tests
+# ---------------------------------------------------------------------------
+
+
+def test_list_models_finds_files_in_subdirectories(tmp_path):
+    model_dir = _make_model_tree(tmp_path)
+    result = list_models(model_dir)
+    names = {f.name for f in result}
+    assert "sd15.safetensors" in names
+    # Two levels deep: loras/SD1.5/detail.safetensors
+    deep = [f for f in result if f.name == "detail.safetensors"]
+    assert len(deep) == 1
+    assert deep[0].relative_to(model_dir) == pathlib.Path("loras/SD1.5/detail.safetensors")
+
+
+def test_list_models_finds_root_level_files(tmp_path):
+    model_dir = _make_model_tree(tmp_path)
+    result = list_models(model_dir)
+    names = {f.name for f in result}
+    assert "root_model.safetensors" in names
+
+
+def test_list_models_returns_empty_for_missing_directory(tmp_path):
+    assert list_models(tmp_path / "nonexistent") == []
+
+
+def test_list_models_ignores_directories(tmp_path):
+    model_dir = _make_model_tree(tmp_path)
+    result = list_models(model_dir)
+    assert all(f.is_file() for f in result)
+    dir_names = {f.name for f in result}
+    assert "empty_dir" not in dir_names
+    assert "checkpoints" not in dir_names
+
+
+# ---------------------------------------------------------------------------
+# list_command tests
+# ---------------------------------------------------------------------------
+
+runner = typer.testing.CliRunner()
+
+
+def test_list_command_shows_type_column(tmp_path):
+    _make_model_tree(tmp_path)
+    with patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path):
+        result = runner.invoke(app, ["list", "--relative-path", "models"])
+    assert result.exit_code == 0
+    assert "Type" in result.output
+    assert "checkpoints" in result.output
+    # Deeply nested type should show full subdirectory path
+    assert "loras/SD1.5" in result.output
+    # Root-level model should still appear
+    assert "root_model.safetensors" in result.output
+
+
+# ---------------------------------------------------------------------------
+# remove tests
+# ---------------------------------------------------------------------------
+
+
+def test_remove_with_path_traversal_is_rejected(tmp_path):
+    model_dir = tmp_path / "models"
+    model_dir.mkdir()
+    (model_dir / "legit.bin").write_bytes(b"x")
+
+    # Create a file outside the model dir that an attacker might target
+    secret = tmp_path / "secret.txt"
+    secret.write_text("sensitive")
+
+    with patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path):
+        result = runner.invoke(
+            app,
+            ["remove", "--relative-path", "models", "--model-names", "../secret.txt", "--confirm"],
+        )
+    # The secret file must still exist
+    assert secret.exists()
+    assert "Invalid model path" in result.output
+
+
+def test_remove_deletes_model_in_subdirectory(tmp_path):
+    model_dir = _make_model_tree(tmp_path)
+    target = model_dir / "checkpoints" / "sd15.safetensors"
+    assert target.exists()
+
+    with patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path):
+        result = runner.invoke(
+            app,
+            ["remove", "--relative-path", "models", "--model-names", "checkpoints/sd15.safetensors", "--confirm"],
+        )
+    assert result.exit_code == 0
+    assert not target.exists()
+
+
+def test_remove_rejects_directory_name(tmp_path):
+    _make_model_tree(tmp_path)
+
+    with patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path):
+        result = runner.invoke(
+            app,
+            ["remove", "--relative-path", "models", "--model-names", "checkpoints", "--confirm"],
+        )
+    # Directory must still exist, not be unlinked
+    assert (tmp_path / "models" / "checkpoints").is_dir()
+    assert "not found" in result.output
+
+
+def test_remove_deletes_root_level_model(tmp_path):
+    model_dir = _make_model_tree(tmp_path)
+    target = model_dir / "root_model.safetensors"
+    assert target.exists()
+
+    with patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path):
+        result = runner.invoke(
+            app,
+            ["remove", "--relative-path", "models", "--model-names", "root_model.safetensors", "--confirm"],
+        )
+    assert result.exit_code == 0
+    assert not target.exists()
+
+
+# ---------------------------------------------------------------------------
+# URL validation tests (existing, kept for completeness)
+# ---------------------------------------------------------------------------
 
 
 def test_valid_model_url():