feat: support JSON-native union inputs

architecture/02-schema.md

@@ -102,15 +102,15 @@ class Runner(BaseRunner):
 
 The resolver handles local imports relative to the predictor file and project root:
 
-| Import Style                         | File Resolved                                               |
-| ------------------------------------ | ----------------------------------------------------------- |
-| `from output_types import X`         | `<project>/output_types.py`                                 |
-| `from .output_types import X`        | `<predictor-dir>/output_types.py`                           |
-| `from models.output import X`        | `<project>/models/output.py`                                |
-| `from .models.output import X`       | `<predictor-dir>/models/output.py`                          |
-| `from output_types import X as Y`    | `<project>/output_types.py` (alias tracked)                 |
-| `from .output_types import X as Y`   | `<predictor-dir>/output_types.py` (alias tracked)           |
-| `from . import output_types`         | `<predictor-dir>/output_types.py` (module alias tracked)    |
+| Import Style                       | File Resolved                                            |
+| ---------------------------------- | -------------------------------------------------------- |
+| `from output_types import X`       | `<project>/output_types.py`                              |
+| `from .output_types import X`      | `<predictor-dir>/output_types.py`                        |
+| `from models.output import X`      | `<project>/models/output.py`                             |
+| `from .models.output import X`     | `<predictor-dir>/models/output.py`                       |
+| `from output_types import X as Y`  | `<project>/output_types.py` (alias tracked)              |
+| `from .output_types import X as Y` | `<predictor-dir>/output_types.py` (alias tracked)        |
+| `from . import output_types`       | `<predictor-dir>/output_types.py` (module alias tracked) |
 
 **How it distinguishes local from external**: the resolver converts the module path to a filesystem path and checks if the file exists. If `output_types.py` exists in the project directory, it's local. If not (e.g., `from transformers import ...`), it's external. Known external packages (stdlib, torch, numpy, etc.) are skipped without a filesystem check.
 
@@ -175,18 +175,28 @@ Each `SchemaType` produces its JSON Schema fragment via `JSONSchema()`:
 
 ### Input Types
 
-| Python                                | JSON Schema                                                      | Notes                      |
-| ------------------------------------- | ---------------------------------------------------------------- | -------------------------- |
-| `str`                                 | `{"type": "string"}`                                             |                            |
-| `int`                                 | `{"type": "integer"}`                                            |                            |
-| `float`                               | `{"type": "number"}`                                             |                            |
-| `bool`                                | `{"type": "boolean"}`                                            |                            |
-| `cog.Path`                            | `{"type": "string", "format": "uri"}`                            | URLs downloaded at runtime |
-| `cog.File`                            | `{"type": "string", "format": "uri"}`                            | File uploads               |
-| `cog.Secret`                          | `{"type": "string", "format": "password", "x-cog-secret": true}` | Masked in logs             |
-| `list[T]`                             | `{"type": "array", "items": {...}}`                              |                            |
-| `Optional[T]`                         | Type T + not in `required`                                       | Input fields only          |
-| `Literal["a", "b"]` / `choices=[...]` | `{"enum": ["a", "b"]}`                                           |                            |
+| Python                                | JSON Schema                                                      | Notes                                                                 |
+| ------------------------------------- | ---------------------------------------------------------------- | --------------------------------------------------------------------- |
+| `str`                                 | `{"type": "string"}`                                             |                                                                       |
+| `int`                                 | `{"type": "integer"}`                                            |                                                                       |
+| `float`                               | `{"type": "number"}`                                             |                                                                       |
+| `bool`                                | `{"type": "boolean"}`                                            |                                                                       |
+| `cog.Path`                            | `{"type": "string", "format": "uri"}`                            | URLs downloaded at runtime                                            |
+| `cog.File`                            | `{"type": "string", "format": "uri"}`                            | File uploads                                                          |
+| `cog.Secret`                          | `{"type": "string", "format": "password", "x-cog-secret": true}` | Masked in logs                                                        |
+| `list[T]`                             | `{"type": "array", "items": {...}}`                              |                                                                       |
+| `Optional[T]` / `T \| None`           | Type T + `nullable: true`, not in `required`                     | Input fields only; never required                                     |
+| `A \| B` / `Union[A, B]`              | `{"anyOf": [A, B]}`                                              | Input-only, JSON-native unions only                                   |
+| `A \| B \| None`                      | `{"anyOf": [A, B]}` + `nullable: true`                           | Multi-variant union; stays in `required` unless a default is supplied |
+| `Literal["a", "b"]` / `choices=[...]` | `{"enum": ["a", "b"]}`                                           |                                                                       |
+
+Input unions are intentionally narrower than output types. Cog supports JSON-native input unions (`str`, `int`, `float`, `bool`, `dict`/`Any`, `list[T]`, and `None`) so request validation can happen at the HTTP boundary and Python normalisation can choose a deterministic value type. Cog rejects unions involving `Path`, `File`, `Secret`, custom coders, and `BaseModel` because those cases are ambiguous for clients or runtime coercion. Output unions remain unsupported (see below).
+
+A plain single-type optional (`Optional[T]` or `T | None`) is **never** placed in `required`, regardless of whether a default is supplied. A multi-variant nullable union (`A | B | None`) is different: because the field carries a concrete `anyOf` value type, it stays in `required` unless a default makes it omittable. This is why the two rows above differ in their `required` behaviour.
+
+Nullable behaviour matches every other optional field: `nullable: true` (plus omission from `required`) means an **omitted** value falls back to the default. An **explicit** JSON `null` is still validated against the field type and is rejected at the HTTP edge, because the runtime validator does not treat OpenAPI's `nullable` keyword as an additional accepted value. "May be null" therefore means "may be omitted", not "accepts an explicit null payload".
+
+> **Runtime caveat:** Cog marks optionals as not-`required` in the schema, but the predictor still needs a Python-level default so the omitted value resolves to `None`. Use `value: Optional[T] = Input(...)` (the `Input(...)` supplies an implicit `None`) or `Input(default=None)`. A bare `value: Optional[T]` annotation with no `= Input(...)` generates a correct "optional" schema but raises `TypeError: missing 1 required positional argument` when the field is omitted at runtime.
 
 ### Output Types
 

docs/python.md

@@ -620,6 +620,30 @@ def run(self, prompt: Optional[str] = Input(description="prompt")) -> str:
 > [!NOTE]
 > `Optional[T]` is supported in `BaseModel` output fields but **not** as a top-level return type. Use a `BaseModel` with optional fields instead.
 
+#### `Union`
+
+Use `A | B` or `Union[A, B]` to accept more than one type for a single input. Cog supports JSON-native union members: `str`, `int`, `float`, `bool`, `dict`/`Any`, `list[T]`, and `None`.
+
+```python
+from cog import BaseRunner, Input
+
+class Runner(BaseRunner):
+    def run(self,
+        value: str | float = Input(description="A string or a number"),
+    ) -> str:
+        return f"{type(value).__name__}:{value}"
+```
+
+At runtime, Cog validates the request against the union and passes the value through as the matching type. For overlapping numeric types, Cog prefers the most specific match (e.g. `bool` before `int`, `int` before `float`), and a JSON integer is accepted for a `float` member.
+
+Combine a union with `None` to make it nullable:
+
+```python
+def run(self, value: str | float | None = Input(default=None)) -> str: ...
+```
+
+Union inputs are validated at the HTTP boundary, so unions involving `Path`, `File`, `Secret`, custom coders, and `BaseModel` are **not** supported, and the build fails if you use them. Union return types are also unsupported — use a `BaseModel` output instead.
+
 #### `list`
 
 Use `list[T]` or `List[T]` to accept or return a list of values. `T` can be a supported Cog type, but nested container types are not supported.

integration-tests/tests/union_input_cli.txtar

@@ -0,0 +1,45 @@
+# Test schema-directed CLI parsing for JSON-native union inputs.
+#
+# value:   str | float  (string member first)
+# flipped: float | str  (number member first)
+#
+#   - "hello" parses as a string because no numeric parse succeeds
+#   - "1.5" parses as a float because the union accepts a number
+#   - "1" parses as an integer (still a valid JSON number for the union)
+#
+# The `flipped` field exercises the case where resolveSchemaType resolves a
+# union to its numeric member first: a non-numeric value must still fall back
+# to the string member instead of erroring.
+#
+# Note: the worker does not coerce primitives at runtime (validation happens
+# at the HTTP edge against the OpenAPI schema), so the CLI must choose the
+# wire type. A bare integer stays an integer; only fractional values become
+# floats. This matches how a plain `float` input also receives a Python int
+# for `-i num=10`.
+
+cog build -t $TEST_IMAGE
+
+cog predict $TEST_IMAGE -i value=hello -i flipped=world
+stdout 'value=str:hello flipped=str:world'
+
+cog predict $TEST_IMAGE -i value=1.5 -i flipped=2.5
+stdout 'value=float:1.5 flipped=float:2.5'
+
+cog predict $TEST_IMAGE -i value=1 -i flipped=2
+stdout 'value=int:1 flipped=int:2'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from cog import BasePredictor
+
+
+class Predictor(BasePredictor):
+    def predict(self, value: str | float, flipped: float | str) -> str:
+        return (
+            f"value={type(value).__name__}:{value} "
+            f"flipped={type(flipped).__name__}:{flipped}"
+        )

integration-tests/tests/union_input_http.txtar

@@ -0,0 +1,54 @@
+# Test JSON-native union inputs over cog serve.
+# value: str | float | None = Input(default=None)
+#   - string accepted, returns "str:..."
+#   - float accepted, returns "float:..."
+#   - integer accepted (valid JSON number), returns "int:..."
+#   - omitted optional defaults to None, returns "NoneType:none"
+#   - explicit null is rejected, matching how every optional field behaves
+#     today (validation happens at the HTTP edge and the runtime validator
+#     does not accept explicit JSON null for a typed field)
+#   - bool rejected (not a member of str | float), returns a validation error
+
+cog build -t $TEST_IMAGE
+
+cog serve
+
+# String member
+curl POST /predictions '{"input":{"value":"hello"}}'
+stdout '"output":"str:hello"'
+
+# Float member
+curl POST /predictions '{"input":{"value":1.5}}'
+stdout '"output":"float:1.5"'
+
+# Integer is a valid JSON number for the union; passed through as int
+curl POST /predictions '{"input":{"value":1}}'
+stdout '"output":"int:1"'
+
+# Omitted optional value defaults to None
+curl POST /predictions '{"input":{}}'
+stdout '"output":"NoneType:none"'
+
+# Explicit null is rejected, consistent with all optional fields
+! curl POST /predictions '{"input":{"value":null}}'
+
+# bool is not a member of str | float -> rejected
+! curl POST /predictions '{"input":{"value":true}}'
+
+# nested object is not a member of str | float -> rejected
+! curl POST /predictions '{"input":{"value":{"x":1}}}'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from cog import BasePredictor, Input
+
+
+class Predictor(BasePredictor):
+    def predict(self, value: str | float | None = Input(default=None)) -> str:
+        if value is None:
+            return "NoneType:none"
+        return f"{type(value).__name__}:{value}"

integration-tests/tests/union_input_list_http.txtar

@@ -0,0 +1,48 @@
+# Test list JSON-native union inputs over cog serve.
+#
+# nums: list[int] | list[float]   (required list union)
+#
+#   - list[int] | list[float] accepts [1] and [1.5] and the empty list []
+#   - list element types are validated: ["3"] and [true] are rejected
+#   - integer elements stay int, fractional elements are float (no runtime
+#     coercion: the wire type is preserved)
+
+cog build -t $TEST_IMAGE
+
+cog serve
+
+# Integer list element kept as int
+curl POST /predictions '{"input":{"nums":[1]}}'
+stdout '"output":"int:1"'
+
+# Float list element
+curl POST /predictions '{"input":{"nums":[1.5]}}'
+stdout '"output":"float:1.5"'
+
+# Empty list is accepted
+curl POST /predictions '{"input":{"nums":[]}}'
+stdout '"output":"empty"'
+
+# String element is not valid for list[int] | list[float] -> rejected
+! curl POST /predictions '{"input":{"nums":["3"]}}'
+
+# bool element is not valid for list[int] | list[float] -> rejected
+! curl POST /predictions '{"input":{"nums":[true]}}'
+
+# A bare scalar is not a list -> rejected
+! curl POST /predictions '{"input":{"nums":1}}'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from cog import BasePredictor
+
+
+class Predictor(BasePredictor):
+    def predict(self, nums: list[int] | list[float]) -> str:
+        if not nums:
+            return "empty"
+        return f"{type(nums[0]).__name__}:{nums[0]}"