fix(cli): validate structure command inputs (#1329)

* fix(cli): validate structure command inputs

Require the prepare-structure model to immediately follow the subcommand so family-option values cannot be mistaken for a model during bootstrap parsing.

Reject equivalent structure and metadata destinations before the legacy runtime writes either output.

Signed-off-by: Vivian Chen <140748220+xuanzic@users.noreply.github.com>

* fix(cli): validate all sample output paths

Precompute and normalize every structure and metadata destination before writing so cross-sample path collisions cannot overwrite earlier outputs.

Cover a two-sample collision and verify that existing destination contents remain unchanged.

Signed-off-by: Vivian Chen <140748220+xuanzic@users.noreply.github.com>

---------

Signed-off-by: Vivian Chen <140748220+xuanzic@users.noreply.github.com>

Author: xuanzic

apps/cli/cli.cpp

@@ -937,12 +937,32 @@ int dispatch(const Command& command, ITask& task, std::ostream& output) {
         const fs::path metadata_path = has_option(command, "--output-json")
                                            ? command.options.at("--output-json")
                                            : structure_path.string() + ".metadata.json";
+        if (fs::absolute(structure_path).lexically_normal() ==
+            fs::absolute(metadata_path).lexically_normal())
+            throw std::invalid_argument("--output and --output-json must use different paths");
         auto indexed_path = [](const fs::path& path, std::size_t index) {
             if (index == 0)
                 return path;
             return path.parent_path() / (path.stem().string() + "_sample_" + std::to_string(index) +
                                          path.extension().string());
         };
+        const auto sample_count = result.samples.empty() ? 1U : result.samples.size();
+        std::vector<std::pair<fs::path, fs::path>> output_paths;
+        output_paths.reserve(sample_count);
+        std::unordered_set<std::string> normalized_paths;
+        for (std::size_t index = 0; index < sample_count; ++index) {
+            auto current_structure = indexed_path(structure_path, index);
+            auto current_metadata = has_option(command, "--output-json")
+                                        ? indexed_path(metadata_path, index)
+                                        : fs::path(current_structure.string() + ".metadata.json");
+            for (const auto& path : {current_structure, current_metadata}) {
+                const auto normalized = fs::absolute(path).lexically_normal().string();
+                if (!normalized_paths.insert(normalized).second)
+                    throw std::invalid_argument(
+                        "--output and --output-json must use different paths");
+            }
+            output_paths.emplace_back(std::move(current_structure), std::move(current_metadata));
+        }
         auto write_output = [](const fs::path& path, const std::string& payload,
                                const char* label) {
             if (!path.parent_path().empty())
@@ -957,11 +977,7 @@ int dispatch(const Command& command, ITask& task, std::ostream& output) {
         auto append_sample = [&](std::size_t index, const std::string& structure,
                                  const std::string& metadata,
                                  const StructureConfidence& confidence) {
-            const auto current_structure = indexed_path(structure_path, index);
-            const auto current_metadata =
-                has_option(command, "--output-json")
-                    ? indexed_path(metadata_path, index)
-                    : fs::path(current_structure.string() + ".metadata.json");
+            const auto& [current_structure, current_metadata] = output_paths.at(index);
             write_output(current_structure, structure, "structure output");
             write_output(current_metadata, metadata, "structure metadata");
             sample_outputs.push_back({{"structure_path", current_structure.string()},

apps/cli/tests/test_sdk_detection_structure_cli.cpp

@@ -119,12 +119,20 @@ void detection(const std::filesystem::path& root) {
 class ExistingStructure final : public trtmc::IStructurePrediction {
   public:
     trtmc::StructurePredictionRequest seen;
+    bool multiple_samples{false};
     trtmc::StructurePredictionResult
     predict_structure(const trtmc::StructurePredictionRequest& request) override {
         seen = request;
         trtmc::StructurePredictionResult result;
         result.structure = "data_existing\n";
         result.metadata_json = "{\"existing\":true}";
+        if (multiple_samples) {
+            result.samples.resize(2);
+            result.samples[0].structure = "data_sample_0\n";
+            result.samples[0].metadata_json = "{\"sample\":0}";
+            result.samples[1].structure = "data_sample_1\n";
+            result.samples[1].metadata_json = "{\"sample\":1}";
+        }
         return result;
     }
 };
@@ -249,6 +257,36 @@ void structure(const std::filesystem::path& root) {
             read_file("existing.cif") == "data_existing\n" &&
             read_file("existing.json") == "{\"existing\":true}",
         "unmigrated structure command preserves every existing explicit argument and binary input");
+    write_file("existing-collision.cif", "keep legacy bytes");
+    auto collision_command =
+        parse({"trtmc", "predict-structure", "unused.bundle", "--input", prepared.string(),
+               "--output", "existing-collision.cif", "--output-json", "./existing-collision.cif"});
+    bool collision_rejected = false;
+    try {
+        (void)trtmc::cli::dispatch(collision_command, existing, existing_output);
+    } catch (const std::invalid_argument& error) {
+        collision_rejected =
+            std::string_view(error.what()).find("different paths") != std::string_view::npos;
+    }
+    check(collision_rejected && read_file("existing-collision.cif") == "keep legacy bytes",
+          "unmigrated structure command rejects equivalent output paths before writing");
+    existing.multiple_samples = true;
+    write_file("cross-sample.cif", "keep metadata bytes");
+    write_file("cross-sample_sample_1.cif", "keep structure bytes");
+    collision_command =
+        parse({"trtmc", "predict-structure", "unused.bundle", "--input", prepared.string(),
+               "--output", "cross-sample_sample_1.cif", "--output-json", "cross-sample.cif"});
+    collision_rejected = false;
+    try {
+        (void)trtmc::cli::dispatch(collision_command, existing, existing_output);
+    } catch (const std::invalid_argument& error) {
+        collision_rejected =
+            std::string_view(error.what()).find("different paths") != std::string_view::npos;
+    }
+    check(collision_rejected && read_file("cross-sample.cif") == "keep metadata bytes" &&
+              read_file("cross-sample_sample_1.cif") == "keep structure bytes",
+          "all multi-sample output paths are validated before writing");
+    existing.multiple_samples = false;
     command.options["--input-encoding"] = "b2rq";
     bool rejected = false;
     try {

core/builder/tensorrt_model_connect/build_cli.py

@@ -7,6 +7,7 @@
 
 import argparse
 import json
+import sys
 from pathlib import Path
 from typing import Sequence
 
@@ -51,12 +52,20 @@ def _parser(prepare_family: object | None = None) -> argparse.ArgumentParser:
 
 
 def main(argv: Sequence[str] | None = None) -> int:
+    arguments = list(sys.argv[1:] if argv is None else argv)
     base_parser = _parser()
-    preliminary, _ = base_parser.parse_known_args(argv)
+    if (
+        len(arguments) > 1
+        and arguments[0] == "prepare-structure"
+        and arguments[1] not in {"-h", "--help"}
+        and arguments[1].startswith("-")
+    ):
+        base_parser.error("MODEL must immediately follow prepare-structure")
+    preliminary, _ = base_parser.parse_known_args(arguments)
     model_dir = _resolve_model(preliminary.model, preliminary.revision)
     family, support = resolve_family(load_model_metadata(model_dir))
     family_module = _load_family(family) if preliminary.command == "prepare-structure" else None
-    args = _parser(family_module).parse_args(argv)
+    args = _parser(family_module).parse_args(arguments)
     if args.command == "prepare-structure":
         prepare = getattr(family_module, "prepare_structure_request", None)
         if not callable(prepare):

core/builder/tests/test_build_cli.py

@@ -366,6 +366,31 @@ def cli_options(args):
     }
 
 
+def test_prepare_structure_requires_model_before_family_options(monkeypatch, capsys) -> None:
+    monkeypatch.setattr(
+        build_cli,
+        "_resolve_model",
+        lambda *_: pytest.fail("ambiguous family option value reached model resolution"),
+    )
+
+    with pytest.raises(SystemExit) as error:
+        build_cli.main(
+            [
+                "prepare-structure",
+                "--num-steps",
+                "300",
+                "model",
+                "--input",
+                "request.yaml",
+                "--output",
+                "request.b2rq",
+            ]
+        )
+
+    assert error.value.code == 2
+    assert "MODEL must immediately follow prepare-structure" in capsys.readouterr().err
+
+
 def test_prepare_structure_uses_the_family_hook_after_task_migration(
     monkeypatch, tmp_path: Path, capsys
 ) -> None: