fix: support TypedDict schema generation

integration-tests/tests/typeddict_input.txtar

@@ -0,0 +1,37 @@
+# Test TypedDict inputs through the default static schema path.
+#
+# This exercises the Go tree-sitter parser end-to-end for TypedDict inputs:
+# build-time extraction, embedded OpenAPI schema, and predict-time JSON input.
+
+# Build the image.
+cog build -t $TEST_IMAGE
+
+# TypedDict and optional list[TypedDict] inputs should be represented like dicts.
+exec docker inspect $TEST_IMAGE --format '{{index .Config.Labels "run.cog.openapi_schema"}}'
+stdout '"data":{"title":"Data","type":"object","x-order":0}'
+stdout '"items":{"default":null,"items":{"type":"object"},"nullable":true,"title":"Items","type":"array","x-order":1}'
+
+# Predict should accept TypedDict-shaped JSON payloads end-to-end.
+cog predict $TEST_IMAGE --json '{"data": {"name": "alice", "count": 2}, "items": [{"name": "first", "count": 1}, {"name": "second", "count": 3}]}'
+stdout '"status": "succeeded"'
+stdout '"output": "alice:2 items=2"'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from typing import TypedDict
+
+from cog import BasePredictor
+
+
+class Payload(TypedDict):
+    name: str
+    count: int
+
+
+class Predictor(BasePredictor):
+    def predict(self, data: Payload, items: list[Payload] | None = None) -> str:
+        return f"{data['name']}:{data['count']} items={len(items or [])}"

integration-tests/tests/typeddict_input_legacy.txtar

@@ -0,0 +1,41 @@
+# Test TypedDict inputs through the legacy runtime schema path.
+#
+# This exercises python/cog/_adt.py via in-container runtime schema generation,
+# proving TypedDict subclasses are treated like dict inputs when
+# COG_LEGACY_SCHEMA=1 forces the legacy path.
+
+env COG_LEGACY_SCHEMA=1
+
+# Build the image through the legacy runtime schema generator.
+cog build -t $TEST_IMAGE
+stderr 'Validating model schema'
+
+# TypedDict and list[TypedDict] inputs work via JSON.
+cog predict $TEST_IMAGE --json '{"data": {"name": "alice", "count": 2}, "items": [{"name": "first", "count": 1}, {"name": "second", "count": 3}]}'
+stdout '"status": "succeeded"'
+stdout '"output": "alice:2 items=2"'
+
+# The generated schema should still expose these as opaque JSON objects.
+exec docker inspect $TEST_IMAGE --format '{{index .Config.Labels "run.cog.openapi_schema"}}'
+stdout '"data":{"title":"Data","type":"object","x-order":0}'
+stdout '"items":{"items":{"type":"object"},"title":"Items","type":"array","x-order":1}'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from typing import TypedDict
+
+from cog import BasePredictor
+
+
+class Payload(TypedDict):
+    name: str
+    count: int
+
+
+class Predictor(BasePredictor):
+    def predict(self, data: Payload, items: list[Payload]) -> str:
+        return f"{data['name']}:{data['count']} items={len(items)}"

integration-tests/tests/typeddict_output.txtar

@@ -0,0 +1,39 @@
+# Test TypedDict return types through the static schema path.
+#
+# This exercises the Go tree-sitter parser end-to-end: build-time extraction,
+# embedded OpenAPI schema, and prediction output for a TypedDict response.
+
+# Build the image.
+cog build -t $TEST_IMAGE
+
+# Verify the schema is embedded with object fields from the TypedDict.
+exec docker inspect $TEST_IMAGE --format '{{index .Config.Labels "run.cog.openapi_schema"}}'
+stdout '"type":"object"'
+stdout '"name":'
+stdout '"count":'
+stdout '"required":\["name"\]'
+
+# Predict should return a JSON object matching the TypedDict.
+cog predict $TEST_IMAGE
+stdout '"name": "alice"'
+stdout '"count": 3'
+
+-- cog.yaml --
+build:
+  python_version: "3.12"
+predict: "predict.py:Predictor"
+
+-- predict.py --
+from typing import NotRequired, TypedDict
+
+from cog import BasePredictor
+
+
+class Output(TypedDict):
+    name: str
+    count: NotRequired[int]
+
+
+class Predictor(BasePredictor):
+    def predict(self) -> Output:
+        return {"name": "alice", "count": 3}