fix!updating preprocess_schemas to create operation schemas for extensions

preprocess_schemas.py

@@ -408,21 +408,26 @@ class name like 'Checkout'); fall back to dot-replaced-with-underscore
 
 def get_required_ops(schema):
     """
-    Scans a schema for the custom 'ucp_request' metadata.
+    Scans a schema's properties and $defs for custom 'ucp_request' metadata.
     Returns a set of operation keys (e.g. {'create', 'update'}) that need distinct models.
     """
     ops = set()
-    properties = schema.get("properties", {})
-    if not isinstance(properties, dict):
-        return ops
-
-    for data in properties.values():
-        if isinstance(data, dict):
-            marker = data.get("ucp_request")
-            if isinstance(marker, str):
-                ops.update(["create", "update"])  # Standard shortcut
-            elif isinstance(marker, dict):
-                ops.update(marker.keys())
+    containers = []
+    if isinstance(schema.get("properties"), dict):
+        containers.append(schema["properties"])
+    if isinstance(schema.get("$defs"), dict):
+        containers.append(schema["$defs"])
+
+    for container in containers:
+        for node in iter_nodes(container):
+            if not isinstance(node, dict):
+                continue
+            marker = node.get("ucp_request")
+            if marker is not None:
+                if isinstance(marker, str):
+                    ops.update(["create", "update"])  # Standard shortcut
+                elif isinstance(marker, dict):
+                    ops.update(marker.keys())
     return ops
 
 
@@ -558,6 +563,20 @@ def _create_single_variant(
             variant, op, file_path, global_variant_requirements
         )
 
+    # Apply request rules to top-level definitions in $defs
+    defs = variant.get("$defs", {})
+    if isinstance(defs, dict):
+        for node in defs.values():
+            if isinstance(node, dict) and (
+                "properties" in node or node.get("type") == "object"
+            ):
+                _apply_request_rules_to_object(
+                    node, op, file_path, global_variant_requirements
+                )
+
+    # Rewrite all external references across the entire variant tree
+    rewrite_refs_to_variants(variant, op, file_path, global_variant_requirements)
+
     return variant
 
 
@@ -626,20 +645,23 @@ def normalize_metadata_schemas(schemas, target_dir):
 
 
 def extract_external_refs(schema, path):
-    """Finds all relative external file references in the schema properties."""
+    """Finds all relative external file references in properties and $defs."""
     refs = []
-    props = schema.get("properties", {})
-    if not isinstance(props, dict):
-        return refs
-
-    for name, data in props.items():
-        for node in iter_nodes(data):
-            if isinstance(node, dict) and "$ref" in node:
-                ref = node["$ref"]
-                ref_file, _, _ = ref.partition("#")
-                if ref_file:
-                    abs_path = str((path.parent / ref_file).resolve())
-                    refs.append((name, abs_path))
+    containers = []
+    if isinstance(schema.get("properties"), dict):
+        containers.append(schema["properties"])
+    if isinstance(schema.get("$defs"), dict):
+        containers.append(schema["$defs"])
+
+    for container in containers:
+        for name, data in container.items():
+            for node in iter_nodes(data):
+                if isinstance(node, dict) and "$ref" in node:
+                    ref = node["$ref"]
+                    ref_file, _, _ = ref.partition("#")
+                    if ref_file:
+                        abs_path = str((path.parent / ref_file).resolve())
+                        refs.append((name, abs_path))
     return refs
 
 
@@ -660,9 +682,10 @@ def propagate_needs_transitive(variant_needs, schema_refs, schemas):
                     if child_path not in schemas:
                         continue
 
-                    # Only propagate if the property isn't 'omit'ted for this op
+                    # Check properties first, then $defs for any operation override
                     data = (
-                        schemas[path].get("properties", {}).get(prop_name, {})
+                        schemas[path].get("properties", {}).get(prop_name)
+                        or schemas[path].get("$defs", {}).get(prop_name, {})
                     )
                     include, _ = eval_prop_inclusion(
                         prop_name, data, op, schemas[path].get("required", [])
@@ -760,3 +783,4 @@ def main():
 
 if __name__ == "__main__":
     main()
+

tests/test_codegen_pipeline.py

@@ -315,6 +315,31 @@ def test_get_required_ops_collects_all_declared_operations(self) -> None:
             {"create", "update", "complete"},
         )
 
+    def test_get_required_ops_collects_from_nested_defs(self) -> None:
+        """Markers inside $defs contribute their operations even if root has none."""
+        schema = {
+            "title": "Extension",
+            "type": "object",
+            "$defs": {
+                "custom": {
+                    "type": "object",
+                    "properties": {
+                        "secret": {"ucp_request": "omit"},
+                        "action": {
+                            "ucp_request": {
+                                "complete": "optional",
+                            }
+                        },
+                    },
+                }
+            },
+        }
+
+        self.assertEqual(
+            preprocess_schemas.get_required_ops(schema),
+            {"create", "update", "complete"},
+        )
+
     def test_eval_prop_inclusion_applies_operation_overrides(self) -> None:
         """Operation markers override base required and inclusion rules."""
         cases = [
@@ -553,6 +578,94 @@ def test_generate_variants_writes_operation_specific_files(self) -> None:
         self.assertEqual(set(update_variant["properties"]), {"id"})
         self.assertEqual(update_variant["required"], ["id"])
 
+    def test_nested_defs_filtered_and_refs_rewritten(self) -> None:
+        """Nested $defs properties are filtered and external refs rewritten."""
+        schema = {
+            "$id": "https://ucp.dev/schemas/food/cart.json",
+            "title": "Cart",
+            "type": "object",
+            "$defs": {
+                "checkout": {
+                    "title": "Checkout with Cart",
+                    "type": "object",
+                    "properties": {
+                        "cart_id": {
+                            "type": "string",
+                            "ucp_request": {
+                                "create": "optional",
+                                "update": "omit",
+                                "complete": "omit",
+                            },
+                        }
+                    },
+                    "allOf": [{"$ref": "checkout.json"}],
+                }
+            },
+            "properties": {
+                "restaurant": {"$ref": "restaurant.json"},
+            },
+        }
+        file_path = Path("/schemas/food/cart.json")
+        checkout_path = str((file_path.parent / "checkout.json").resolve())
+        restaurant_path = str((file_path.parent / "restaurant.json").resolve())
+        variant_needs = {
+            checkout_path: {"create", "update", "complete"},
+            restaurant_path: {"create", "update"},
+        }
+
+        create_variant = preprocess_schemas._create_single_variant(
+            schema,
+            "create",
+            "cart",
+            file_path,
+            variant_needs,
+        )
+        update_variant = preprocess_schemas._create_single_variant(
+            schema,
+            "update",
+            "cart",
+            file_path,
+            variant_needs,
+        )
+        complete_variant = preprocess_schemas._create_single_variant(
+            schema,
+            "complete",
+            "cart",
+            file_path,
+            variant_needs,
+        )
+
+        # Create variant includes cart_id (without ucp_request) and rewrites checkout.json
+        checkout_def_create = create_variant["$defs"]["checkout"]
+        self.assertIn("cart_id", checkout_def_create["properties"])
+        self.assertNotIn(
+            "ucp_request", checkout_def_create["properties"]["cart_id"]
+        )
+        self.assertEqual(
+            checkout_def_create["allOf"][0]["$ref"],
+            "checkout_create_request.json",
+        )
+        self.assertEqual(
+            create_variant["properties"]["restaurant"]["$ref"],
+            "restaurant_create_request.json",
+        )
+
+        # Update variant omits cart_id and rewrites checkout.json
+        checkout_def_update = update_variant["$defs"]["checkout"]
+        self.assertEqual(checkout_def_update["properties"], {})
+        self.assertEqual(
+            checkout_def_update["allOf"][0]["$ref"],
+            "checkout_update_request.json",
+        )
+
+        # Complete variant omits cart_id and rewrites checkout.json
+        checkout_def_complete = complete_variant["$defs"]["checkout"]
+        self.assertEqual(checkout_def_complete["properties"], {})
+        self.assertEqual(
+            checkout_def_complete["allOf"][0]["$ref"],
+            "checkout_complete_request.json",
+        )
+
 
 class PipelineDependencyTest(unittest.TestCase):
     """Tests metadata normalization and transitive variant dependencies."""
@@ -814,6 +927,125 @@ def test_propagation_with_fragment(self) -> None:
                 "child_create_request.json#/$defs/item",
             )
 
+    def test_main_preprocesses_nested_capability_extensions_end_to_end(
+        self,
+    ) -> None:
+        """Capability extensions in $defs trigger variants and propagate refs."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            root = Path(temp_dir)
+            preprocess_schemas.save_json(
+                {
+                    "$defs": {
+                        "entity": {
+                            "type": "object",
+                            "properties": {"id": {"type": "string"}},
+                            "required": ["id"],
+                        }
+                    }
+                },
+                root / "ucp.json",
+            )
+            preprocess_schemas.save_json(
+                {
+                    "$id": "https://ucp.dev/schemas/checkout.json",
+                    "title": "Checkout",
+                    "type": "object",
+                    "properties": {
+                        "id": {
+                            "type": "string",
+                            "ucp_request": {
+                                "create": "omit",
+                                "update": "required",
+                            },
+                        }
+                    },
+                },
+                root / "checkout.json",
+            )
+            preprocess_schemas.save_json(
+                {
+                    "$id": "https://ucp.dev/schemas/cart.json",
+                    "title": "Cart",
+                    "type": "object",
+                    "$defs": {
+                        "checkout": {
+                            "title": "Checkout with Cart",
+                            "type": "object",
+                            "properties": {
+                                "cart_id": {
+                                    "type": "string",
+                                    "ucp_request": {
+                                        "create": "optional",
+                                        "update": "omit",
+                                    },
+                                }
+                            },
+                            "allOf": [{"$ref": "checkout.json"}],
+                        }
+                    },
+                    "properties": {
+                        "items": {"type": "array"},
+                    },
+                },
+                root / "cart.json",
+            )
+
+            with (
+                mock.patch.object(
+                    sys,
+                    "argv",
+                    ["preprocess_schemas.py", str(root)],
+                ),
+                contextlib.redirect_stdout(io.StringIO()),
+            ):
+                preprocess_schemas.main()
+
+            self.assertTrue(
+                (root / "checkout_create_request.json").exists(),
+                "checkout_create_request.json was not generated",
+            )
+            self.assertTrue(
+                (root / "checkout_update_request.json").exists(),
+                "checkout_update_request.json was not generated",
+            )
+            self.assertTrue(
+                (root / "cart_create_request.json").exists(),
+                "cart_create_request.json was not generated",
+            )
+            self.assertTrue(
+                (root / "cart_update_request.json").exists(),
+                "cart_update_request.json was not generated",
+            )
+
+            cart_create = preprocess_schemas.load_json(
+                root / "cart_create_request.json"
+            )
+            cart_update = preprocess_schemas.load_json(
+                root / "cart_update_request.json"
+            )
+
+            # Check cart_create_request.json
+            self.assertIn(
+                "cart_id", cart_create["$defs"]["checkout"]["properties"]
+            )
+            self.assertNotIn(
+                "ucp_request",
+                cart_create["$defs"]["checkout"]["properties"]["cart_id"],
+            )
+            self.assertEqual(
+                cart_create["$defs"]["checkout"]["allOf"][0]["$ref"],
+                "checkout_create_request.json",
+            )
+
+            # Check cart_update_request.json
+            self.assertEqual(
+                cart_update["$defs"]["checkout"]["properties"], {}
+            )
+            self.assertEqual(
+                cart_update["$defs"]["checkout"]["allOf"][0]["$ref"],
+                "checkout_update_request.json",
+            )
+
 
 class MetadataUnionTest(unittest.TestCase):
     """The UcpMetadata root union is derived from ucp.json $defs."""
@@ -1977,3 +2209,4 @@ def test_payment_handler_base_rejects_invalid_version(self) -> None:
 
 if __name__ == "__main__":
     unittest.main()
+