fix: restore payload_builder wrapping for dict schemas

Author: abhishekmadan30

src/taskgraph/transforms/task.py

@@ -16,11 +16,14 @@
 from dataclasses import dataclass
 from typing import Callable, Literal, Optional, Union
 
+import voluptuous
+
 from taskgraph.transforms.base import TransformSequence
 from taskgraph.util.hash import hash_path
 from taskgraph.util.keyed_by import evaluate_keyed_by
 from taskgraph.util.schema import (
     IndexSchema,
+    LegacySchema,
     OptimizationType,
     Schema,
     TaskPriority,
@@ -194,6 +197,14 @@ class PayloadBuilder:
 
 
 def payload_builder(name, schema):
+    if isinstance(schema, dict):
+        schema = LegacySchema(
+            {
+                voluptuous.Required("implementation"): name,
+                voluptuous.Optional("os"): str,
+            }
+        ).extend(schema)
+
     def wrap(func):
         assert name not in payload_builders, f"duplicate payload builder name {name}"
         payload_builders[name] = PayloadBuilder(schema, func)

src/taskgraph/util/schema.py

@@ -48,6 +48,9 @@ def validate_schema(schema, obj, msg_prefix):
         # Handle plain Python types (e.g. str, int) via msgspec.convert
         elif isinstance(schema, type):
             msgspec.convert(obj, schema)
+        # Handle plain dict schemas (e.g. from downstream payload builders)
+        elif isinstance(schema, dict):
+            voluptuous.Schema(schema)(obj)
         else:
             raise TypeError(f"Unsupported schema type: {type(schema)}")
     except (

test/test_transforms_task.py

@@ -2,11 +2,13 @@
 Tests for the 'task' transforms.
 """
 
+import unittest
 from copy import deepcopy
 from pathlib import Path
 from pprint import pprint
 
 import pytest
+import voluptuous
 from pytest_taskgraph import FakeParameters
 
 from taskgraph.transforms import task
@@ -966,3 +968,38 @@ def test_task_priority(run_transform, graph_config, test_task):
         assert task_dict["task"]["priority"] == priority
     else:
         assert task_dict["task"]["priority"] == graph_config["task-priority"]
+
+
+class TestPayloadBuilderDictSchema(unittest.TestCase):
+    """Regression test: payload_builder must wrap dict schemas with
+    'implementation' and 'os' keys via LegacySchema.extend(), matching
+    the v18.1 behavior."""
+
+    def test_dict_schema_wraps_implementation_key(self):
+        """A dict schema should be extended to require 'implementation'."""
+        # Clear any previously-registered test builder
+        task.payload_builders.pop("test-builder", None)
+
+        @task.payload_builder("test-builder", schema={"command": [str]})
+        def _builder(config, task, task_def):
+            pass
+
+        schema = task.payload_builders["test-builder"].schema
+        assert isinstance(schema, voluptuous.Schema)
+
+        # Valid: includes implementation matching the builder name
+        schema({"implementation": "test-builder", "command": ["echo"]})
+
+        # Valid: os is optional
+        schema({"implementation": "test-builder", "command": ["echo"], "os": "linux"})
+
+        # Invalid: wrong implementation value
+        with self.assertRaises(voluptuous.MultipleInvalid):
+            schema({"implementation": "wrong-name", "command": ["echo"]})
+
+        # Invalid: missing implementation
+        with self.assertRaises(voluptuous.MultipleInvalid):
+            schema({"command": ["echo"]})
+
+        # Cleanup
+        task.payload_builders.pop("test-builder", None)

test/test_util_schema.py

@@ -269,6 +269,16 @@ def test_index_schema_accepts_all_fields(self):
             },
             "index schema",
         )
+class TestValidateSchemaDictHandler(unittest.TestCase):
+    """validate_schema must accept plain dict schemas passed
+    by downstream payload builders without raising TypeError."""
+
+    def test_dict_schema_valid(self):
+        validate_schema({"name": str, "count": int}, {"name": "a", "count": 1}, "pfx")
+
+    def test_dict_schema_invalid(self):
+        with self.assertRaises(Exception):
+            validate_schema({"name": str}, {"name": 123}, "pfx")
 
 
 def test_optionally_keyed_by():