Fix discriminated union schema gen bug (#8904)

pydantic/_internal/_discriminated_union.py

@@ -54,6 +54,12 @@ def inner(s: core_schema.CoreSchema, recurse: _core_utils.Recurse) -> core_schem
         if discriminator is not None:
             if definitions is None:
                 definitions = collect_definitions(schema)
+                # After we collect the definitions schemas, we must run through the discriminator
+                # application logic for each one. This step is crucial to prevent an exponential
+                # increase in complexity that occurs if schemas are left as 'union' schemas
+                # rather than 'tagged-union' schemas.
+                # For more details, see https://github.com/pydantic/pydantic/pull/8904#discussion_r1504687302
+                definitions = {k: recurse(v, inner) for k, v in definitions.items()}
             s = apply_discriminator(s, discriminator, definitions)
         return s
 

tests/test_discriminated_union.py

@@ -1852,3 +1852,27 @@ class SubModel(MyModel):
         'title': 'MyModel',
         'type': 'object',
     }
+
+
+def test_nested_schema_gen_uses_tagged_union_in_ref() -> None:
+    class NestedState(BaseModel):
+        state_type: Literal['nested']
+        substate: 'AnyState'
+
+    # If this type is left out, the model behaves normally again
+    class LoopState(BaseModel):
+        state_type: Literal['loop']
+        substate: 'AnyState'
+
+    class LeafState(BaseModel):
+        state_type: Literal['leaf']
+
+    AnyState = Annotated[Union[NestedState, LoopState, LeafState], Field(..., discriminator='state_type')]
+    NestedState.model_rebuild()
+    LoopState.model_rebuild()
+    adapter = TypeAdapter(AnyState)
+
+    assert adapter.core_schema['schema']['type'] == 'tagged-union'
+    for definition in adapter.core_schema['definitions']:
+        if definition['schema']['model_name'] in ['NestedState', 'LoopState']:
+            assert definition['schema']['fields']['substate']['schema']['schema']['type'] == 'tagged-union'