feat: model input validation before build and run

architecture/02-schema.md

@@ -328,12 +328,12 @@ A simplified example showing a multi-file predictor with structured output:
 
 ## Code References
 
-| File                        | Purpose                                                              |
-| --------------------------- | -------------------------------------------------------------------- |
-| `pkg/schema/schema_type.go` | `SchemaType` ADT, `ResolveSchemaType()`, `JSONSchema()` generation   |
-| `pkg/schema/types.go`       | `PredictorInfo`, `PrimitiveType`, `FieldType`, `InputField`, imports |
-| `pkg/schema/python/`        | Tree-sitter Python parser and cross-file resolution                  |
-| `pkg/schema/openapi.go`     | OpenAPI document assembly from `PredictorInfo`                       |
-| `pkg/schema/generator.go`   | Top-level `Generate()`, `GenerateCombined()`, `Parser` type          |
-| `pkg/schema/errors.go`      | Typed schema error kinds                                             |
-| `pkg/image/build.go`        | Build-time schema generation entry point and schema file validation  |
+| File                         | Purpose                                                              |
+| ---------------------------- | -------------------------------------------------------------------- |
+| `pkg/schema/schema_type.go`  | `SchemaType` ADT, `ResolveSchemaType()`, `JSONSchema()` generation   |
+| `pkg/schema/types.go`        | `PredictorInfo`, `PrimitiveType`, `FieldType`, `InputField`, imports |
+| `pkg/schema/python/`         | Tree-sitter Python parser and cross-file resolution                  |
+| `pkg/schema/openapi_spec.go` | OpenAPI document assembly from `PredictorInfo`                       |
+| `pkg/schema/generator.go`    | Top-level `Generate()`, `GenerateCombined()`, `Parser` type          |
+| `pkg/schema/errors.go`       | Typed schema error kinds                                             |
+| `pkg/image/build.go`         | Build-time schema generation entry point and schema file validation  |

architecture/05-build-system.md

@@ -12,9 +12,10 @@ flowchart TB
         weights["weights"]
     end
 
-    subgraph cli["CLI (pkg/cli/build.go)"]
+    subgraph cli["Build Orchestration (pkg/image/)"]
         parse["Parse Config"]
         validate["Validate"]
+        schema["Generate OpenAPI Schema"]
     end
 
     subgraph generate["Dockerfile Generation (pkg/dockerfile/)"]
@@ -30,12 +31,12 @@ flowchart TB
     end
 
     subgraph post["Post-Build"]
-        schema["Generate OpenAPI Schema"]
         freeze["pip freeze"]
         labels["Apply Labels"]
     end
 
-    yaml --> parse --> validate
+    yaml --> parse --> validate --> schema
+    code --> schema
     validate --> generator
     compat --> generator
     baseimage --> generator
@@ -44,7 +45,6 @@ flowchart TB
     code --> buildkit
     weights --> buildkit
     buildkit --> image
-    image --> schema
     image --> freeze
     schema --> labels
     freeze --> labels

architecture/06-cli.md

@@ -45,11 +45,12 @@ cog run -i prompt="A photo of a cat" -i steps=50
 
 What happens:
 
-1. Builds the image (if needed)
-2. Starts a container running the [Container Runtime](./04-container-runtime.md)
-3. Parses `-i` flags against the [Schema](./02-schema.md)
-4. Sends a [PredictionRequest](./03-prediction-api.md) to the container's HTTP API
-5. Streams output back to terminal
+1. Resolves the [Schema](./02-schema.md) from local source or an existing image label
+2. Parses, coerces, and validates `-i` flags when that schema is available
+3. Builds the image (for local source)
+4. Starts a container running the [Container Runtime](./04-container-runtime.md)
+5. Falls back to the runtime schema for older images without a usable schema label
+6. Sends a [PredictionRequest](./03-prediction-api.md) and streams output to the terminal
 
 Input types are inferred from the schema:
 
@@ -112,9 +113,9 @@ What happens (see [Build System](./05-build-system.md) for details):
 
 1. **Parse** `cog.yaml`
 2. **Resolve** CUDA/cuDNN versions from compatibility matrix
-3. **Generate** Dockerfile
-4. **Build** image via Docker/Buildkit
-5. **Run** container to extract [Schema](./02-schema.md)
+3. **Generate** the [Schema](./02-schema.md) statically from model source
+4. **Generate** Dockerfile
+5. **Build** image via Docker/Buildkit
 6. **Apply** labels (schema, config, pip freeze)
 
 Key flags:
@@ -174,15 +175,18 @@ There's also a separate `base-image` binary (`cmd/base-image/`) with subcommands
 
 ## How CLI Commands Interact with Containers
 
-Commands like `predict` and `serve` follow the same pattern: build an image, start a container, communicate via HTTP. The CLI never runs model code directly.
+Local-source predictions generate and validate against the schema before building, then start a container and communicate over HTTP. Existing-image predictions use the image's schema label when available and otherwise fall back to the runtime schema after startup. `cog serve` builds and starts the same runtime without parsing prediction inputs. The CLI never runs model code directly.
+
+The local-source prediction flow is:
 
 ```mermaid
 sequenceDiagram
     participant CLI as cog CLI (host)
     participant Docker
     participant Container as Container (runtime)
 
-    CLI->>CLI: Parse -i flags, load cog.yaml
+    CLI->>CLI: Load cog.yaml and generate schema
+    CLI->>CLI: Parse, coerce, and validate -i flags
     CLI->>Docker: Build image (if needed)
     Docker-->>CLI: Image ready
 
@@ -234,7 +238,7 @@ Commands delegate to packages under `pkg/`:
 - `pkg/dockerfile/` -- Dockerfile generation and base image selection
 - `pkg/docker/` -- Docker client operations
 - `pkg/predict/` -- Local prediction execution (talks to container's HTTP API)
-- `pkg/schema/` -- Static schema generator (tree-sitter, experimental)
+- `pkg/schema/` -- Static schema generator (tree-sitter)
 - `pkg/wheels/` -- SDK and coglet wheel resolution
 
 **Infrastructure:**

integration-tests/tests/input_validation_before_build.txtar

@@ -0,0 +1,102 @@
+# Local-source inputs are validated against the statically-generated schema
+# *before* the Docker image is built. Every command below must fail at preflight
+# with a clear message and must not start a build. Covers run, the deprecated
+# predict alias, and train, across -i and --json, for the full range of
+# validation errors (unknown/missing/type/range/length/regex/enum/null).
+
+# --- Unknown input name (-i and --json) ---
+! cog run -i prompt=hi -i nope=1
+stderr 'unknown input "nope"'
+! stderr 'Building Docker image'
+
+! cog run --json '{"prompt": "hi", "nope": 1}'
+stderr 'unknown input "nope"'
+! stderr 'Building Docker image'
+
+# --- Missing a required input ---
+! cog run
+stderr 'missing required input "prompt"'
+! stderr 'Building Docker image'
+
+# --- Numeric bounds: ge (0 < 1) and le (99 > 10) ---
+! cog run -i prompt=hi -i count=0
+stderr 'invalid input "count"'
+! stderr 'Building Docker image'
+
+! cog run -i prompt=hi -i count=99
+stderr 'invalid input "count"'
+! stderr 'Building Docker image'
+
+# --- Type mismatch (not an integer) ---
+! cog run -i prompt=hi -i count=lots
+stderr 'invalid input "count": expected integer, got string'
+! stderr 'Building Docker image'
+
+# --- Choices / enum ---
+! cog run -i prompt=hi -i size=huge
+stderr 'invalid input "size": must be one of: small, medium, large'
+! stderr 'Building Docker image'
+
+# --- String length (max_length=5) ---
+! cog run -i prompt=hi -i title=toolong
+stderr 'invalid input "title"'
+! stderr 'Building Docker image'
+
+# --- Regex (must match ^[a-z]+$) ---
+! cog run -i prompt=hi -i slug=BAD
+stderr 'invalid input "slug"'
+! stderr 'Building Docker image'
+
+# --- Explicit null is rejected, even for an optional/nullable field ---
+! cog run --json '{"prompt": "hi", "count": null}'
+stderr 'must not be null'
+! stderr 'Building Docker image'
+
+! cog run --json '{"prompt": "hi", "note": null}'
+stderr 'must not be null'
+! stderr 'Building Docker image'
+
+# --- The deprecated `cog predict` local path validates the same way ---
+! cog predict -i prompt=hi -i count=0
+stderr 'invalid input "count"'
+! stderr 'Building Docker image'
+
+# --- Training inputs are validated before build too ---
+! cog train -i epochs=0
+stderr 'invalid input "epochs"'
+! stderr 'Building Docker image'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+train: "train.py:train"
+
+-- predict.py --
+from typing import Optional
+
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(
+        self,
+        prompt: str,
+        count: int = Input(default=5, ge=1, le=10),
+        size: str = Input(default="small", choices=["small", "medium", "large"]),
+        title: str = Input(default="ok", max_length=5),
+        slug: str = Input(default="ok", regex="^[a-z]+$"),
+        note: Optional[str] = Input(default=None),
+    ) -> str:
+        return prompt
+
+-- train.py --
+from cog import BaseModel, Input
+
+
+class TrainingOutput(BaseModel):
+    weights: str
+
+
+def train(epochs: int = Input(ge=1)) -> TrainingOutput:
+    return TrainingOutput(weights="ok")

integration-tests/tests/input_validation_before_start.txtar

@@ -0,0 +1,32 @@
+# Inputs to a pre-built image are validated against the image's embedded schema
+# *before* the container starts. Builds once, then exercises the validation
+# paths (unknown/missing) via both -i and --json. None may start the container.
+
+cog build -t $TEST_IMAGE
+
+# --- Unknown input name (-i and --json) ---
+! cog predict $TEST_IMAGE -i nope=value
+stderr 'unknown input "nope"'
+! stderr 'Starting Docker image and running setup'
+
+! cog predict $TEST_IMAGE --json '{"nope": 1}'
+stderr 'unknown input "nope"'
+! stderr 'Starting Docker image and running setup'
+
+# --- Missing a required input ---
+! cog predict $TEST_IMAGE
+stderr 'missing required input "prompt"'
+! stderr 'Starting Docker image and running setup'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from cog import BasePredictor
+
+
+class Predictor(BasePredictor):
+    def predict(self, prompt: str) -> str:
+        return prompt

integration-tests/tests/string_predictor.txtar

@@ -6,9 +6,14 @@ cog build -t $TEST_IMAGE
 cog predict $TEST_IMAGE -i s=world
 stdout 'hello world'
 
-# Missing required input fails
+# Unknown input fails
 ! cog predict $TEST_IMAGE -i wrong=value
-stderr 's: Field required'
+stderr 'unknown input "wrong"'
+
+# Omitting a required input fails before the container starts
+! cog predict $TEST_IMAGE
+stderr 'missing required input "s"'
+! stderr 'Starting Docker image and running setup'
 
 -- cog.yaml --
 build: