Repository: pydantic/pydantic
Pull request: Include JSON Schema input core schema in function schemas
<!-- Thank you for your contribution! -->
<!-- Unless your change is trivial, please create an issue to discuss the change before creating a PR -->

## Change Summary

Partly fixes https://github.com/pydantic/pydantic/issues/11033.
Requires https://github.com/pydantic/pydantic-core/pull/1572

<!-- Please give a short summary of the changes. -->

## Related issue number

<!-- WARNING: please use "fix #123" style references so the issue is closed when this PR is merged. -->

## Checklist

* [ ] The pull request title is a good summary of the changes - it will be used in the changelog
* [ ] Unit tests for the changes exist
* [ ] Tests pass on CI
* [ ] Documentation reflects the changes where applicable
* [ ] My PR is ready to review, **please add a comment including the phrase "please review" to assign reviewers**


Problem statement:
JSON Schema gen chokes on `BeforeValidator` with a union in the `json_schema_input_type`
### Initial Checks

- [X] I confirm that I'm using Pydantic V2

### Description

Using a `BeforeValidator` with a union type (see `example.py` below) results in :

```
Traceback (most recent call last):
...
  File "/<...>/.venv/lib/python3.11/site-packages/pydantic/json_schema.py", line 2245, in _add_json_refs
    _add_json_refs(self.definitions[defs_ref])
                   ~~~~~~~~~~~~~~~~^^^^^^^^^^
KeyError: '__main____TypeA-Input__1'
```

### Example Code

```Python
from pydantic.plugin import TypeAlias
from typing_extensions import TypeAliasType
from typing import Annotated
from pydantic import BaseModel, BeforeValidator, AfterValidator

class TypeA(BaseModel):
    pass

class TypeB(BaseModel):
    pass

FieldType = Annotated[list[str], BeforeValidator(lambda x: x if isinstance(x, list) else [x], TypeA | TypeB)]

class Model(BaseModel):
    field: FieldType

print(Model.model_json_schema())
```


### Python, Pydantic & OS Version

```Text
pydantic version: 2.10.0b2
        pydantic-core version: 2.27.0
          pydantic-core build: profile=release pgo=false
                 install path: /Users/joshcannon/code/shimbboleth/.venv/lib/python3.11/site-packages/pydantic
               python version: 3.11.6 | packaged by conda-forge | (main, Oct  3 2023, 10:37:07) [Clang 15.0.7 ]
                     platform: macOS-14.7.1-arm64-arm-64bit
             related packages: typing_extensions-4.12.2
                       commit: unknown
```

Hints from issue discussion:
I suspect this may be related to https://github.com/pydantic/pydantic/blob/aee6057378ccfec02126bf9c984a9b6d6b411777/pydantic/functional_validators.py#L130

(maybe it doesn't handle more complex schemas?)
cc @Viicos
Oh yeah the `input_schema` is `{'type': 'union', 'choices': [{'type': 'definition-ref', 'schema_ref': '__main__.TypeA:5517702256'}, {'type': 'definition-ref', 'schema_ref': '__main__.TypeB:5805303600'}]}`

so the `resolve_ref_schema` call should be recursive over the input schema
Which means the workaround is obviously to make the input schema.... a `TypeAliasType`! 😂 
I think the easiest here would be to just have `json_schema_input_type` core schema as part of the validator core schema (so that the definition is not removed, because anything in the metadata is ignored). Or else, as @thejcannon mentioned we need to recursively inline references, but this means we need to make use of the `WalkCoreSchema` utils which presumably are going to be removed if we go along with Markus' work. @sydney-runkle wdyt?
@Viicos, if I understand your point correctly, we'd be migrating this info out of metadata and onto the core schema directly?

I'm in favor of this approach for now, you know how I dislike our current metadata pattern 😆 

If you'd like, I'm happy to take this task on.

Commit message: Merge branch 'main' into js-input-type-in-cs

Patch to review:
diff --git a/pydantic/_internal/_core_metadata.py b/pydantic/_internal/_core_metadata.py
index c8d452443cb..89e3e78833c 100644
--- a/pydantic/_internal/_core_metadata.py
+++ b/pydantic/_internal/_core_metadata.py
@@ -4,8 +4,6 @@
 from warnings import warn
 
 if TYPE_CHECKING:
-    from pydantic_core import CoreSchema
-
     from ..config import JsonDict, JsonSchemaExtraCallable
     from ._schema_generation_shared import (
         GetJsonSchemaFunction,
@@ -20,7 +18,6 @@ class CoreMetadata(TypedDict, total=False):
         pydantic_js_annotation_functions: List of JSON schema functions that don't resolve refs during application.
         pydantic_js_prefer_positional_arguments: Whether JSON schema generator will
             prefer positional over keyword arguments for an 'arguments' schema.
-        pydantic_js_input_core_schema: Schema associated with the input value for the associated
             custom validation function. Only applies to before, plain, and wrap validators.
         pydantic_js_udpates: key / value pair updates to apply to the JSON schema for a type.
         pydantic_js_extra: WIP, either key/value pair updates to apply to the JSON schema, or a custom callable.
@@ -37,7 +34,6 @@ class CoreMetadata(TypedDict, total=False):
     pydantic_js_functions: list[GetJsonSchemaFunction]
     pydantic_js_annotation_functions: list[GetJsonSchemaFunction]
     pydantic_js_prefer_positional_arguments: bool
-    pydantic_js_input_core_schema: CoreSchema
     pydantic_js_updates: JsonDict
     pydantic_js_extra: JsonDict | JsonSchemaExtraCallable
 
diff --git a/pydantic/_internal/_core_utils.py b/pydantic/_internal/_core_utils.py
index e4c1aed9c1f..f6ab20e43df 100644
--- a/pydantic/_internal/_core_utils.py
+++ b/pydantic/_internal/_core_utils.py
@@ -280,12 +280,38 @@ def handle_dict_schema(self, schema: core_schema.DictSchema, f: Walk) -> core_sc
             schema['values_schema'] = self.walk(values_schema, f)
         return schema
 
-    def handle_function_schema(self, schema: AnyFunctionSchema, f: Walk) -> core_schema.CoreSchema:
-        if not is_function_with_inner_schema(schema):
-            return schema
+    def handle_function_after_schema(
+        self, schema: core_schema.AfterValidatorFunctionSchema, f: Walk
+    ) -> core_schema.CoreSchema:
         schema['schema'] = self.walk(schema['schema'], f)
         return schema
 
+    def handle_function_before_schema(
+        self, schema: core_schema.BeforeValidatorFunctionSchema, f: Walk
+    ) -> core_schema.CoreSchema:
+        schema['schema'] = self.walk(schema['schema'], f)
+        if 'json_schema_input_schema' in schema:
+            schema['json_schema_input_schema'] = self.walk(schema['json_schema_input_schema'], f)
+        return schema
+
+    # TODO duplicate schema types for serializers and validators, needs to be deduplicated:
+    def handle_function_plain_schema(
+        self, schema: core_schema.PlainValidatorFunctionSchema | core_schema.PlainSerializerFunctionSerSchema, f: Walk
+    ) -> core_schema.CoreSchema:
+        if 'json_schema_input_schema' in schema:
+            schema['json_schema_input_schema'] = self.walk(schema['json_schema_input_schema'], f)
+        return schema  # pyright: ignore[reportReturnType]
+
+    # TODO duplicate schema types for serializers and validators, needs to be deduplicated:
+    def handle_function_wrap_schema(
+        self, schema: core_schema.WrapValidatorFunctionSchema | core_schema.WrapSerializerFunctionSerSchema, f: Walk
+    ) -> core_schema.CoreSchema:
+        if 'schema' in schema:
+            schema['schema'] = self.walk(schema['schema'], f)
+        if 'json_schema_input_schema' in schema:
+            schema['json_schema_input_schema'] = self.walk(schema['json_schema_input_schema'], f)
+        return schema  # pyright: ignore[reportReturnType]
+
     def handle_union_schema(self, schema: core_schema.UnionSchema, f: Walk) -> core_schema.CoreSchema:
         new_choices: list[CoreSchema | tuple[CoreSchema, str]] = []
         for v in schema['choices']:
diff --git a/pydantic/functional_validators.py b/pydantic/functional_validators.py
index 0892e3460d2..aa3da8c2713 100644
--- a/pydantic/functional_validators.py
+++ b/pydantic/functional_validators.py
@@ -127,14 +127,6 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
             if self.json_schema_input_type is PydanticUndefined
             else handler.generate_schema(self.json_schema_input_type)
         )
-        # Try to resolve the original schema if required, because schema cleaning
-        # won't inline references in metadata:
-        if input_schema is not None:
-            try:
-                input_schema = handler.resolve_ref_schema(input_schema)
-            except LookupError:
-                pass
-        metadata = {'pydantic_js_input_core_schema': input_schema} if input_schema is not None else {}
 
         info_arg = _inspect_validator(self.func, 'before')
         if info_arg:
@@ -143,11 +135,13 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
                 func,
                 schema=schema,
                 field_name=handler.field_name,
-                metadata=metadata,
+                json_schema_input_schema=input_schema,
             )
         else:
             func = cast(core_schema.NoInfoValidatorFunction, self.func)
-            return core_schema.no_info_before_validator_function(func, schema=schema, metadata=metadata)
+            return core_schema.no_info_before_validator_function(
+                func, schema=schema, json_schema_input_schema=input_schema
+            )
 
     @classmethod
     def _from_decorator(cls, decorator: _decorators.Decorator[_decorators.FieldValidatorDecoratorInfo]) -> Self:
@@ -229,13 +223,6 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
             serialization = None
 
         input_schema = handler.generate_schema(self.json_schema_input_type)
-        # Try to resolve the original schema if required, because schema cleaning
-        # won't inline references in metadata:
-        try:
-            input_schema = handler.resolve_ref_schema(input_schema)
-        except LookupError:
-            pass
-        metadata = {'pydantic_js_input_core_schema': input_schema} if input_schema is not None else {}
 
         info_arg = _inspect_validator(self.func, 'plain')
         if info_arg:
@@ -244,14 +231,14 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
                 func,
                 field_name=handler.field_name,
                 serialization=serialization,  # pyright: ignore[reportArgumentType]
-                metadata=metadata,
+                json_schema_input_schema=input_schema,
             )
         else:
             func = cast(core_schema.NoInfoValidatorFunction, self.func)
             return core_schema.no_info_plain_validator_function(
                 func,
                 serialization=serialization,  # pyright: ignore[reportArgumentType]
-                metadata=metadata,
+                json_schema_input_schema=input_schema,
             )
 
     @classmethod
@@ -312,14 +299,6 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
             if self.json_schema_input_type is PydanticUndefined
             else handler.generate_schema(self.json_schema_input_type)
         )
-        # Try to resolve the original schema if required, because schema cleaning
-        # won't inline references in metadata:
-        if input_schema is not None:
-            try:
-                input_schema = handler.resolve_ref_schema(input_schema)
-            except LookupError:
-                pass
-        metadata = {'pydantic_js_input_core_schema': input_schema} if input_schema is not None else {}
 
         info_arg = _inspect_validator(self.func, 'wrap')
         if info_arg:
@@ -328,14 +307,14 @@ def __get_pydantic_core_schema__(self, source_type: Any, handler: GetCoreSchemaH
                 func,
                 schema=schema,
                 field_name=handler.field_name,
-                metadata=metadata,
+                json_schema_input_schema=input_schema,
             )
         else:
             func = cast(core_schema.NoInfoWrapValidatorFunction, self.func)
             return core_schema.no_info_wrap_validator_function(
                 func,
                 schema=schema,
-                metadata=metadata,
+                json_schema_input_schema=input_schema,
             )
 
     @classmethod
diff --git a/pydantic/json_schema.py b/pydantic/json_schema.py
index 301332e9ba7..3471e0df589 100644
--- a/pydantic/json_schema.py
+++ b/pydantic/json_schema.py
@@ -1053,9 +1053,7 @@ def function_before_schema(self, schema: core_schema.BeforeValidatorFunctionSche
         Returns:
             The generated JSON schema.
         """
-        if self._mode == 'validation' and (
-            input_schema := schema.get('metadata', {}).get('pydantic_js_input_core_schema')
-        ):
+        if self._mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')):
             return self.generate_inner(input_schema)
 
         return self.generate_inner(schema['schema'])
@@ -1080,9 +1078,7 @@ def function_plain_schema(self, schema: core_schema.PlainValidatorFunctionSchema
         Returns:
             The generated JSON schema.
         """
-        if self._mode == 'validation' and (
-            input_schema := schema.get('metadata', {}).get('pydantic_js_input_core_schema')
-        ):
+        if self._mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')):
             return self.generate_inner(input_schema)
 
         return self.handle_invalid_for_json_schema(
@@ -1098,9 +1094,7 @@ def function_wrap_schema(self, schema: core_schema.WrapValidatorFunctionSchema)
         Returns:
             The generated JSON schema.
         """
-        if self._mode == 'validation' and (
-            input_schema := schema.get('metadata', {}).get('pydantic_js_input_core_schema')
-        ):
+        if self._mode == 'validation' and (input_schema := schema.get('json_schema_input_schema')):
             return self.generate_inner(input_schema)
 
         return self.generate_inner(schema['schema'])
diff --git a/pyproject.toml b/pyproject.toml
index 01f65ec5126..d1dcd3cebad 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -176,6 +176,9 @@ markers = [
 [tool.uv]
 default-groups = ['dev']
 
+[tool.uv.sources]
+pydantic-core = { git = "https://github.com/pydantic/pydantic-core", branch = "js-input-type-in-cs" }
+
 # configuring https://github.com/pydantic/hooky
 [tool.hooky]
 reviewers = ['sydney-runkle']
diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py
index 40f826ef3a5..c0c4a6394b1 100644
--- a/tests/test_json_schema.py
+++ b/tests/test_json_schema.py
@@ -6542,31 +6542,39 @@ def validate_f(cls, value: Any) -> int: ...
 @pytest.mark.parametrize(
     'validator',
     [
-        PlainValidator(lambda v: v, json_schema_input_type='Sub'),
-        BeforeValidator(lambda v: v, json_schema_input_type='Sub'),
-        WrapValidator(lambda v, h: h(v), json_schema_input_type='Sub'),
+        PlainValidator(lambda v: v, json_schema_input_type='Union[Sub1, Sub2]'),
+        BeforeValidator(lambda v: v, json_schema_input_type='Union[Sub1, Sub2]'),
+        WrapValidator(lambda v, h: h(v), json_schema_input_type='Union[Sub1, Sub2]'),
     ],
 )
 def test_json_schema_input_type_with_refs(validator) -> None:
-    """Test that `'definition-ref` schemas for `json_schema_input_type` are inlined.
+    """Test that `'definition-ref` schemas for `json_schema_input_type` are supported.
 
     See: https://github.com/pydantic/pydantic/issues/10434.
+    See: https://github.com/pydantic/pydantic/issues/11033
     """
 
-    class Sub(BaseModel):
+    class Sub1(BaseModel):
+        pass
+
+    class Sub2(BaseModel):
         pass
 
     class Model(BaseModel):
         sub: Annotated[
-            Sub,
-            PlainSerializer(lambda v: v, return_type=Sub),
+            Union[Sub1, Sub2],
+            PlainSerializer(lambda v: v, return_type=Union[Sub1, Sub2]),
             validator,
         ]
 
     json_schema = Model.model_json_schema()
 
-    assert 'Sub' in json_schema['$defs']
-    assert json_schema['properties']['sub']['$ref'] == '#/$defs/Sub'
+    assert 'Sub1' in json_schema['$defs']
+    assert 'Sub2' in json_schema['$defs']
+    assert json_schema['properties']['sub'] == {
+        'anyOf': [{'$ref': '#/$defs/Sub1'}, {'$ref': '#/$defs/Sub2'}],
+        'title': 'Sub',
+    }
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 820fe9d322c..ba838264688 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -637,29 +637,6 @@ def walk(s, recurse):
     }
 
 
-def test_handle_function_schema():
-    schema = core_schema.with_info_before_validator_function(
-        lambda v, _info: v, core_schema.float_schema(), field_name='field_name'
-    )
-
-    def walk(s, recurse):
-        # change type to str
-        if s['type'] == 'float':
-            s['type'] = 'str'
-        return s
-
-    schema = _WalkCoreSchema().handle_function_schema(schema, walk)
-    assert schema['type'] == 'function-before'
-    assert schema['schema'] == {'type': 'str'}
-
-    def walk1(s, recurse):
-        # this is here to make sure this function is not called
-        assert False
-
-    schema = _WalkCoreSchema().handle_function_schema(core_schema.int_schema(), walk1)
-    assert schema['type'] == 'int'
-
-
 def test_handle_call_schema():
     param_a = core_schema.arguments_parameter(name='a', schema=core_schema.str_schema(), mode='positional_only')
     args_schema = core_schema.arguments_schema([param_a])


Code review comments:
diff_hunk: @@ -176,6 +176,9 @@ markers = [
 [tool.uv]
 default-groups = ['dev']
 
+[tool.uv.sources]
+pydantic-core = { git = "https://github.com/pydantic/pydantic-core", branch = "js-input-type-in-cs" }
+
path: pyproject.toml
text: @user1:
```suggestion
```