Let console-codegen take several spec files

composer/input/parsing.py

@@ -1,4 +1,5 @@
 import argparse
+import pathlib
 from typing import TypeVar, Protocol, cast, Annotated, get_type_hints, get_origin, Any, get_args, Union
 from composer.input.types import CommandLineArgs, ResumeArgs, Arg, OptionalArg, RAGDBOptions, ModelOptions, LanggraphOptions, UploadPaths, InputData, SpecInput
 from composer.input.files import DOCUMENT_SUFFIXES, FileUploader
@@ -142,7 +143,11 @@ def _common_options(parser: argparse.ArgumentParser) -> None:
 def fresh_workflow_argument_parser() -> TypedArgumentParser[CommandLineArgs]:
     """Configure command line argument parser."""
     parser = argparse.ArgumentParser(description="Certora AI Composer for Smart Contract Generation")
-    parser.add_argument("spec_file", help="Specification file for the smart contract")
+    parser.add_argument(
+        "spec_file", nargs="+",
+        help="One or more specification files for the smart contract. All of them "
+             "gate the generated code: it must satisfy every rule in every spec.",
+    )
     parser.add_argument("interface_file", help="The interface file for the smart contract")
     parser.add_argument("system_doc", help="A text document describing the system")
     _common_options(parser)
@@ -153,25 +158,52 @@ def fresh_workflow_argument_parser() -> TypedArgumentParser[CommandLineArgs]:
 async def upload_input(uploader: FileUploader, args: UploadPaths) -> InputData:
     """Turn the CLI's spec / interface / system-doc paths into an ``InputData``.
 
-    Spec and interface are unconditionally uploaded to the Files API as text
+    Specs and interface are unconditionally uploaded to the Files API as text
     (``upload_text_file_if_needed`` → ``UploadedTextFile``, a ``TextDocument``);
     the system doc goes through ``get_document`` so a PDF is uploaded while a
     text design doc stays inline.
+
+    Each spec is materialized in the VFS under its own file name, since
+    ``vfs_path`` keys the specs downstream (audit's resume artifact indexes them
+    by it). A single spec keeps the conventional ``rules.spec`` name so existing
+    single-spec runs and their recorded artifacts are unaffected.
     """
-    spec = await uploader.upload_text_file_if_needed(args.spec_file)
+    specs = [
+        SpecInput(file=await uploader.upload_text_file_if_needed(path), vfs_path=vfs_path)
+        for path, vfs_path in zip(args.spec_file, _spec_vfs_paths(args.spec_file))
+    ]
     intf = await uploader.upload_text_file_if_needed(args.interface_file)
     system_doc = await uploader.get_document(args.system_doc)
     if system_doc is None:
         raise FileNotFoundError(f"System document not found or not a file: {args.system_doc}")
-    # The legacy CLI triad is single-spec; map it to a one-element specs list at
-    # the conventional codegen path. The pipeline is plumbed for N specs.
     return InputData(
-        specs=[SpecInput(file=spec, vfs_path="rules.spec")],
+        specs=specs,
         system_doc=system_doc,
         intf=intf,
     )
 
 
+def _spec_vfs_paths(spec_files: list[str]) -> list[str]:
+    """VFS names for the CLI's spec paths: the file names, deduplicated by path.
+
+    Distinct directories can hold same-named specs (``core/vault.spec`` and
+    ``periphery/vault.spec``), and a collision would silently drop one of them
+    from anything keyed by ``vfs_path``, so reject it with the offending name
+    rather than inventing a suffix the user never asked for.
+    """
+    if len(spec_files) == 1:
+        return ["rules.spec"]
+    names = [pathlib.PurePath(p).name for p in spec_files]
+    duplicates = {n for n in names if names.count(n) > 1}
+    if duplicates:
+        raise ValueError(
+            f"Spec file names must be unique; got {len(names)} specs with repeated "
+            f"name(s): {', '.join(sorted(duplicates))}. Rename or copy them so each "
+            f"spec has a distinct file name."
+        )
+    return names
+
+
 def _common_resume_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("--commentary", default=None, help="Commentary describing the changes to the system. If prefixed with @, assumed to be a filename from which the commentary is read")
     parser.add_argument("src_thread_id", help="The thread id from which to resume the workflow")

composer/input/types.py

@@ -141,7 +141,9 @@ class ExtendedModelOptions(_ModelOptionsCommon, Protocol):
     )]
 
 class UploadPaths(Protocol):
-    spec_file: str
+    # One or more specs, all gating the same generated contract (argparse
+    # ``nargs="+"``), hence a list even for the single-spec invocation.
+    spec_file: list[str]
     interface_file: str
     system_doc: str
 

tests/test_codegen_multi_spec_input.py

@@ -0,0 +1,88 @@
+"""``console-codegen`` accepts N specs, all gating the same generated contract.
+
+The workflow has been plumbed for several specs (``InputData.specs``), but the
+CLI mapped its triad to a one-element list, so there was no way to hand it more
+than one. Natspec emits one spec per component — four for a modest contract —
+and merging them by hand means reconciling four copies of the same ERC20 ghost
+model, so the CLI is the thing that needed to move.
+
+``vfs_path`` keys the specs downstream (audit's resume artifact indexes by it),
+so these pin the naming: one spec keeps the conventional ``rules.spec``, several
+take their file names, and a name collision is refused rather than silently
+dropping a spec.
+"""
+
+import pytest
+
+from composer.input.files import FileUploader
+from composer.input.parsing import fresh_workflow_argument_parser, upload_input
+
+
+class FakeUploader(FileUploader):
+    """Records what it was asked to upload; returns the path as the document."""
+
+    def __init__(self) -> None:
+        self.uploaded: list[str] = []
+
+    async def upload_text_file_if_needed(self, path: str):  # type: ignore[override]
+        self.uploaded.append(path)
+        return path
+
+    async def get_document(self, path):  # type: ignore[override]
+        return str(path)
+
+    async def _upload_bytes(self, crc_basename: str, file_data: bytes, mime: str) -> str:
+        raise AssertionError("upload_input should not reach the binary upload path")
+
+
+def _parse(argv: list[str]):
+    parser = fresh_workflow_argument_parser()
+    import sys
+    old, sys.argv = sys.argv, ["console-codegen", *argv]
+    try:
+        return parser.parse_args()
+    finally:
+        sys.argv = old
+
+
+def test_single_spec_parses_as_before() -> None:
+    args = _parse(["rules.spec", "IFoo.sol", "design.md"])
+
+    assert args.spec_file == ["rules.spec"]
+    assert args.interface_file == "IFoo.sol"
+    assert args.system_doc == "design.md"
+
+
+def test_several_specs_are_collected() -> None:
+    args = _parse(["a.spec", "b.spec", "c.spec", "IFoo.sol", "design.md"])
+
+    assert args.spec_file == ["a.spec", "b.spec", "c.spec"]
+    assert args.interface_file == "IFoo.sol"
+    assert args.system_doc == "design.md"
+
+
+@pytest.mark.asyncio
+async def test_one_spec_keeps_the_conventional_vfs_name() -> None:
+    args = _parse(["some/where/views.spec", "IFoo.sol", "design.md"])
+
+    data = await upload_input(FakeUploader(), args)
+
+    assert [s.vfs_path for s in data.specs] == ["rules.spec"]
+
+
+@pytest.mark.asyncio
+async def test_several_specs_are_named_after_their_files() -> None:
+    args = _parse(["core/views.spec", "core/withdrawal.spec", "IFoo.sol", "design.md"])
+
+    data = await upload_input(FakeUploader(), args)
+
+    assert [s.vfs_path for s in data.specs] == ["views.spec", "withdrawal.spec"]
+    assert [s.file for s in data.specs] == ["core/views.spec", "core/withdrawal.spec"]
+
+
+@pytest.mark.asyncio
+async def test_colliding_spec_names_are_refused() -> None:
+    args = _parse(["core/vault.spec", "periphery/vault.spec", "IFoo.sol", "design.md"])
+
+    with pytest.raises(ValueError, match="vault.spec"):
+        await upload_input(FakeUploader(), args)