From d11d3bb84d284e119d74aa77dddbaca6686b8476 Mon Sep 17 00:00:00 2001 From: Sushant Poudel Date: Sun, 18 Oct 2026 08:51:29 +0545 Subject: [PATCH] fix(models): do not cache a discriminator mapping built from unbuilt schemas Reading a discriminated union caches the key -> variant mapping on the union for the life of the process. `_build_discriminated_union_meta` reads each variant's `__pydantic_core_schema__`, and the SDK defers Pydantic builds, so a variant that another thread is still building has no readable schema yet. That variant was silently skipped and the resulting incomplete mapping was cached anyway, so the union then decoded to the wrong variant until the process restarted, with the first-variant fallback that exists for exactly this kind of mismatch made useless because the discriminator lookup silently pointed at the wrong type. The skip in that loop conflated two different situations, which is why the `None` return from `_extract_field_schema_pv2` could not be reused: a variant whose schema has not finished building cannot contribute, whereas a variant that genuinely does not declare the discriminator legitimately contributes nothing. `_has_built_model_schema` separates them. A variant whose schema is not yet a model/model-fields schema marks the mapping incomplete, and an incomplete mapping is returned to the caller but not cached, so the next call rebuilds it once the build has finished. Complete mappings cache exactly as before. The test is deterministic rather than threaded: it replaces one variant's `__pydantic_core_schema__` with an unbuilt one, which is the state the racing thread produces, so no thread timing is involved. --- src/anthropic/_models.py | 31 ++++++++++++++++++++++++++++++- tests/test_models.py | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/src/anthropic/_models.py b/src/anthropic/_models.py index c8d30222b..723e86b65 100644 --- a/src/anthropic/_models.py +++ b/src/anthropic/_models.py @@ -804,6 +804,11 @@ def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any, mapping: dict[str, type] = {} discriminator_alias: str | None = None + # A variant whose schema has not finished building yet cannot contribute to + # the mapping. Caching the resulting incomplete mapping would leave the union + # decoding to the wrong variant for the rest of the process, so an incomplete + # mapping is returned to the caller but never cached. + complete = True for variant in get_args(union): variant = strip_annotated_type(variant) @@ -811,6 +816,7 @@ def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any, if PYDANTIC_V1: field_info = cast("dict[str, FieldInfo]", variant.__fields__).get(discriminator_field_name) # pyright: ignore[reportDeprecated, reportUnnecessaryCast] if not field_info: + complete = False continue # Note: if one variant defines an alias then they all should @@ -821,8 +827,14 @@ def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any, if isinstance(entry, str): mapping[entry] = variant else: + if not _has_built_model_schema(variant): + # Another thread is still running this model's deferred build. + complete = False + continue + field = _extract_field_schema_pv2(variant, discriminator_field_name) if not field: + # This variant genuinely does not declare the discriminator. continue # Note: if one variant defines an alias then they all should @@ -843,10 +855,27 @@ def _build_discriminated_union_meta(*, union: type, meta_annotations: tuple[Any, discriminator_field=discriminator_field_name, discriminator_alias=discriminator_alias, ) - DISCRIMINATOR_CACHE.setdefault(key, details) + if complete: + DISCRIMINATOR_CACHE.setdefault(key, details) return details +def _has_built_model_schema(model: type[BaseModel]) -> bool: + """Whether `model` has finished its deferred Pydantic build. + + The SDK defers Pydantic builds, so a model that another thread is still + building exposes a schema that is not yet a `model`/`model-fields` schema. + Reading the discriminator field from it is meaningless, and skipping such a + variant would produce an incomplete mapping. + """ + schema = model.__pydantic_core_schema__ + if schema["type"] == "definitions": + schema = schema["schema"] + if schema["type"] != "model": + return False + return schema["schema"]["type"] == "model-fields" + + def _extract_field_schema_pv2(model: type[BaseModel], field_name: str) -> ModelField | None: schema = model.__pydantic_core_schema__ if schema["type"] == "definitions": diff --git a/tests/test_models.py b/tests/test_models.py index 6654ba752..447e36417 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -916,6 +916,41 @@ class B(BaseModel): @pytest.mark.skipif(PYDANTIC_V1, reason="TypeAliasType is not supported in Pydantic v1") +def test_discriminated_unions_incomplete_mapping_is_not_cached() -> None: + class A(BaseModel): + type: Literal["a"] + + data: str + + class B(BaseModel): + type: Literal["b"] + + data: int + + UnionType = cast(Any, Union[A, B]) + annotated = cast(Any, Annotated[UnionType, UnionDiscriminator(field_name="type")]) + + real_schema = B.__pydantic_core_schema__ + try: + # Stand in for the race: a model whose deferred Pydantic build has not + # finished exposes a schema that is not a model schema yet, so its + # discriminator cannot be read. Caching the incomplete mapping would + # leave this union decoding to the wrong variant for the process's life. + B.__pydantic_core_schema__ = {"type": "definition-ref"} # type: ignore[assignment] + + m = construct_type(value={"type": "a", "data": "foo"}, type_=annotated) + assert isinstance(m, A) + assert not DISCRIMINATOR_CACHE, "an incomplete discriminator mapping must not be cached" + finally: + B.__pydantic_core_schema__ = real_schema # type: ignore[assignment] + + # With every schema readable the complete mapping is cached and reused. + construct_type(value={"type": "b", "data": "x"}, type_=annotated) + assert DISCRIMINATOR_CACHE, "the complete mapping should be cached" + cached = next(iter(DISCRIMINATOR_CACHE.values())) + assert set(cached.mapping) == {"a", "b"} + + def test_type_alias_type() -> None: class Model(BaseModel): alias: Alias