openhands-sdk 1.7.3__py3-none-any.whl → 1.7.4__py3-none-any.whl

openhands/sdk/utils/models.py

@@ -1,23 +1,38 @@
 import inspect
-import json
 import logging
-import os
+import threading
 from abc import ABC
-from typing import Annotated, Any, ClassVar, Literal, NoReturn, Self, Union
+from typing import Annotated, Any, Self, Union
 
 from pydantic import (
     BaseModel,
     Discriminator,
-    Field,
+    ModelWrapValidatorHandler,
+    SerializationInfo,
+    SerializerFunctionWrapHandler,
     Tag,
-    TypeAdapter,
-    ValidationError,
+    ValidationInfo,
+    computed_field,
+    model_serializer,
+    model_validator,
 )
-from pydantic_core import ErrorDetails
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import CoreSchema
 
 
 logger = logging.getLogger(__name__)
-_rebuild_required = True
+
+# Thread-local storage for tracking schemas currently being generated.
+# This prevents infinite recursion when generating JSON schemas for
+# discriminated unions that reference each other.
+_thread_local = threading.local()
+
+
+def _get_schemas_in_progress() -> dict[type, JsonSchemaValue]:
+    """Get the thread-local dict for tracking in-progress schema generation."""
+    if not hasattr(_thread_local, "schemas_in_progress"):
+        _thread_local.schemas_in_progress = {}
+    return _thread_local.schemas_in_progress
 
 
 def _is_abstract(type_: type) -> bool:
@@ -28,27 +43,6 @@ def _is_abstract(type_: type) -> bool:
         return False
 
 
-def _get_all_subclasses(cls) -> set[type]:
-    """
-    Recursively finds and returns all (loaded) subclasses of a given class.
-    """
-    result = set()
-    for subclass in cls.__subclasses__():
-        result.add(subclass)
-        result.update(_get_all_subclasses(subclass))
-    return result
-
-
-def rebuild_all():
-    """Rebuild all polymorphic classes."""
-    global _rebuild_required
-    _rebuild_required = False
-    for cls in _get_all_subclasses(OpenHandsModel):
-        cls.model_rebuild(force=True)
-    for cls in _get_all_subclasses(DiscriminatedUnionMixin):
-        cls.model_rebuild(force=True)
-
-
 def kind_of(obj) -> str:
     """Get the string value for the kind tag"""
     if isinstance(obj, dict):
@@ -58,46 +52,15 @@ def kind_of(obj) -> str:
     return obj.__name__
 
 
-def _create_enhanced_discriminated_union_error_message(
-    invalid_kind: str, cls_name: str, valid_kinds: list[str]
-) -> str:
-    """Create an enhanced error message for discriminated union validation failures."""
-    possible_kinds_str = ", ".join(sorted(valid_kinds)) if valid_kinds else "none"
-    return (
-        f"Unexpected kind '{invalid_kind}' for {cls_name}. "
-        f"Expected one of: {possible_kinds_str}. "
-        f"If you receive this error when trying to wrap a "
-        f"DiscriminatedUnion instance inside another pydantic model, "
-        f"you may need to use OpenHandsModel instead of BaseModel "
-        f"to make sure that an invalid schema has not been cached."
-    )
-
-
-def _extract_invalid_kind_from_validation_error(error: ErrorDetails) -> str:
-    """Extract the invalid kind from a Pydantic validation error."""
-    input_value = error.get("input")
-    if input_value is not None and hasattr(input_value, "kind"):
-        return input_value.kind
-    elif isinstance(input_value, dict) and "kind" in input_value:
-        return input_value["kind"]
-    else:
-        return kind_of(input_value)
-
-
-def _handle_discriminated_union_validation_error(
-    validation_error: ValidationError, cls_name: str, valid_kinds: list[str]
-) -> NoReturn:
-    """Handle discriminated union validation errors with enhanced messages."""
-    for error in validation_error.errors():
-        if error.get("type") == "union_tag_invalid":
-            invalid_kind = _extract_invalid_kind_from_validation_error(error)
-            error_msg = _create_enhanced_discriminated_union_error_message(
-                invalid_kind, cls_name, valid_kinds
-            )
-            raise ValueError(error_msg) from validation_error
-
-    # If it's not a discriminated union error, re-raise the original error
-    raise validation_error
+def _get_all_subclasses(cls) -> set[type]:
+    """
+    Recursively finds and returns all (loaded) subclasses of a given class.
+    """
+    result = set()
+    for subclass in cls.__subclasses__():
+        result.add(subclass)
+        result.update(_get_all_subclasses(subclass))
+    return result
 
 
 def get_known_concrete_subclasses(cls) -> list[type]:
@@ -115,176 +78,194 @@ def get_known_concrete_subclasses(cls) -> list[type]:
     return out
 
 
-class OpenHandsModel(BaseModel):
-    """
-    Tags a class where the which may be a discriminated union or contain fields
-    which contain a discriminated union. The first time an instance is initialized,
-    the schema is loaded, or a model is validated after a subclass is defined we
-    regenerate all the polymorphic mappings.
-    """
+def _get_checked_concrete_subclasses(cls: type) -> dict[str, type]:
+    result = {}
+    for sub in get_known_concrete_subclasses(cls):
+        existing = result.get(sub.__name__)
+        if existing:
+            raise ValueError(
+                f"Duplicate class definition for {cls.__module__}.{cls.__name__}: "
+                f"{existing.__module__}.{existing.__name__} : "
+                f"{sub.__module__}.{sub.__name__}"
+            )
+        if "<locals>" in sub.__qualname__:
+            raise ValueError(
+                f"Local classes not supported! {sub.__module__}.{sub.__name__} "
+                f"/ {cls.__module__}.{cls.__name__} "
+                "(Since they may not exist at deserialization time)"
+            )
+        result[sub.__name__] = sub
+    return result
 
-    def model_post_init(self, _context):
-        _rebuild_if_required()
 
-    @classmethod
-    def model_validate(cls, *args, **kwargs) -> Self:
-        _rebuild_if_required()
-        return super().model_validate(*args, **kwargs)
+class OpenHandsModel(BaseModel):
+    """Deprecated: This class exists only for backward compatibility.
 
-    @classmethod
-    def model_validate_json(cls, *args, **kwargs) -> Self:
-        _rebuild_if_required()
-        return super().model_validate_json(*args, **kwargs)
+    This class is no longer required for discriminated union support.
+    New code should extend pydantic.BaseModel directly instead of OpenHandsModel.
 
-    @classmethod
-    def model_json_schema(cls, *args, **kwargs) -> dict[str, Any]:
-        _rebuild_if_required()
-        return super().model_json_schema(*args, **kwargs)
-
-    def model_dump_json(self, **kwargs):
-        # This was overridden because it seems there is a bug where sometimes
-        # duplicate fields are produced by model_dump_json which does not appear
-        # in model_dump
-        kwargs["mode"] = "json"
-        return json.dumps(self.model_dump(**kwargs), ensure_ascii=False)
-
-    def __init_subclass__(cls, **kwargs):
-        """
-        When a new subclass is defined, mark that we will need
-        to rebuild everything
-        """
-        global _rebuild_required
-        _rebuild_required = True
+    Existing code that extends OpenHandsModel will continue to work, but
+    migration to BaseModel is recommended.
+    """
 
-        return super().__init_subclass__(**kwargs)
 
+class DiscriminatedUnionMixin(OpenHandsModel):
+    @computed_field
+    @property
+    def kind(self) -> str:
+        return self.__class__.__name__
 
-class DiscriminatedUnionMixin(OpenHandsModel, ABC):
-    """A Base class for members of tagged unions discriminated by the class name.
+    @model_validator(mode="wrap")
+    @classmethod
+    def _validate_subtype(
+        cls, data: Any, handler: ModelWrapValidatorHandler[Self], info: ValidationInfo
+    ) -> Self:
+        if isinstance(data, cls):
+            return data
+        kind = data.pop("kind", None)
+        if not _is_abstract(cls):
+            # Sanity check: if we're validating a concrete class directly,
+            # the kind (if provided) should match the class name. This should
+            # always be true at this point since resolve_kind() would have
+            # already routed to the correct subclass.
+            assert kind is None or kind == cls.__name__
+            return handler(data)
+        if kind is None:
+            subclasses = _get_checked_concrete_subclasses(cls)
+            if not subclasses:
+                raise ValueError(
+                    f"No kinds defined for {cls.__module__}.{cls.__name__}"
+                )
+            elif len(subclasses) == 1:
+                # If there is ony 1 possible implementation, then we do not need
+                # to state the kind explicitly - it can only be this!
+                kind = next(iter(subclasses))
+            else:
+                # There is more than 1 kind defined but the input did not specify
+                # This will cause an error to be raised
+                kind = ""
+        subclass = cls.resolve_kind(kind)
+        return subclass.model_validate(data, context=info.context)
+
+    @model_serializer(mode="wrap")
+    def _serialize_by_kind(
+        self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
+    ):
+        if isinstance(self, dict):
+            # Sometimes pydantic passes a dict in here.
+            return self
+        if self._is_handler_for_current_class(handler):
+            result = handler(self)
+            return result
+
+        # Delegate to the implementing class
+        result = self.model_dump(
+            mode=info.mode,
+            context=info.context,
+            by_alias=info.by_alias,
+            exclude_unset=info.exclude_unset,
+            exclude_defaults=info.exclude_defaults,
+            exclude_none=info.exclude_none,
+            exclude_computed_fields=info.exclude_computed_fields,
+            round_trip=info.round_trip,
+            serialize_as_any=info.serialize_as_any,
+        )
+        return result
 
-    This class provides automatic subclass registration and discriminated union
-    functionality. Each subclass is automatically registered when defined and
-    can be used for polymorphic serialization/deserialization.
+    def _is_handler_for_current_class(
+        self, handler: SerializerFunctionWrapHandler
+    ) -> bool:
+        """Check if the handler is for this class by parsing its repr string.
 
-    Child classes will automatically have a type field defined, which is used as a
-    discriminator for union types.
-    """
+        WARNING: This is a fragile approach that relies on Pydantic's internal
+        repr format for SerializerFunctionWrapHandler. The handler is a Pydantic
+        wrapper around a Rust function that provides no public API for determining
+        which class it serializes. Parsing the repr string is the only available
+        mechanism.
 
-    __pydantic_core_schema__: ClassVar[Any]
-    __pydantic_validator__: ClassVar[Any]
-    __pydantic_serializer__: ClassVar[Any]
+        Expected format: `SerializationCallable(serializer=<ClassName>)`
 
-    kind: str = Field(default="")  # We dynamically update on a per class basis
+        If Pydantic changes this format, multiple unit tests will fail immediately,
+        including tests in test_discriminated_union.py that verify serialization
+        behavior across the class hierarchy.
+        """
+        # should be in the format `SerializationCallable(serializer=<NAME>)`
+        repr_str = str(handler)
 
-    @classmethod
-    def resolve_kind(cls, kind: str) -> type:
-        for subclass in get_known_concrete_subclasses(cls):
-            if subclass.__name__ == kind:
-                return subclass
-
-        # Generate enhanced error message for unknown kind
-        valid_kinds = [
-            subclass.__name__ for subclass in get_known_concrete_subclasses(cls)
-        ]
-        error_msg = _create_enhanced_discriminated_union_error_message(
-            kind, cls.__name__, valid_kinds
-        )
-        raise ValueError(error_msg)
+        # Get everything after =
+        _, name = repr_str.split("=", 1)
 
-    @classmethod
-    def __get_pydantic_core_schema__(cls, source_type, handler):
-        """Generate discriminated union schema for TypeAdapter compatibility."""
-        if cls.__name__ == "DiscriminatedUnionMixin":
-            return handler(source_type)
-
-        if _is_abstract(source_type):
-            _rebuild_if_required()
-            serializable_type = source_type.get_serializable_type()
-            # If there are subclasses, generate schema for the discriminated union
-            if serializable_type is not source_type:
-                from pydantic_core import core_schema
-
-                # Generate the base schema
-                base_schema = handler.generate_schema(serializable_type)
-
-                # Wrap it with a custom validation function that provides
-                # enhanced error messages
-                def validate_with_enhanced_error(value, handler_func, info):  # noqa: ARG001
-                    try:
-                        return handler_func(value)
-                    except ValidationError as e:
-                        valid_kinds = [
-                            subclass.__name__
-                            for subclass in get_known_concrete_subclasses(source_type)
-                        ]
-                        _handle_discriminated_union_validation_error(
-                            e, source_type.__name__, valid_kinds
-                        )
-
-                # Create a with_info_wrap_validator_function schema
-                return core_schema.with_info_wrap_validator_function(
-                    validate_with_enhanced_error,
-                    base_schema,
-                )
+        # Cut off the )
+        name = name[:-1]
 
-        return handler(source_type)
+        result = self.__class__.__name__ == name
+        return result
 
     @classmethod
-    def __get_pydantic_json_schema__(cls, core_schema, handler):
-        """Add discriminator to OpenAPI schema and ensure component generation."""
-        json_schema = handler(core_schema)
+    def __get_pydantic_json_schema__(
+        cls, core_schema: CoreSchema, handler: Any
+    ) -> JsonSchemaValue:
+        schemas_in_progress = _get_schemas_in_progress()
+
+        # First we check if we are already generating a schema
+        schema = schemas_in_progress.get(cls)
+        if schema:
+            return schema
+
+        # Set a temp schema to prevent infinite recursion
+        schemas_in_progress[cls] = {"$ref": f"#/$defs/{cls.__name__}"}
+        try:
+            if _is_abstract(cls):
+                subclasses = _get_checked_concrete_subclasses(cls)
+                if not subclasses:
+                    raise ValueError(f"No subclasses defined for {cls.__name__}")
+                if len(subclasses) == 1:
+                    # Use the shared generator for single subclass too
+                    gen = handler.generate_json_schema
+                    sub_schema = gen.generate_inner(
+                        next(iter(subclasses.values())).__pydantic_core_schema__
+                    )
+                    return sub_schema
 
-        # Add discriminator if this is a oneOf schema
-        if isinstance(json_schema, dict) and "oneOf" in json_schema:
-            # Add title for abstract classes to encourage separate component creation
-            if _is_abstract(cls) and "title" not in json_schema:
-                json_schema["title"] = cls.__name__
+                # Use the shared generator to properly register definitions
+                gen = handler.generate_json_schema
+                schemas = []
+                for sub in subclasses.values():
+                    sub_schema = gen.generate_inner(sub.__pydantic_core_schema__)
+                    schemas.append(sub_schema)
 
-            if "discriminator" not in json_schema:
+                # Build discriminator mapping from $ref schemas
                 mapping = {}
-                for option in json_schema["oneOf"]:
+                for option in schemas:
                     if "$ref" in option:
                         kind = option["$ref"].split("/")[-1]
                         mapping[kind] = option["$ref"]
 
-                if mapping:
-                    json_schema["discriminator"] = {
-                        "propertyName": "kind",
-                        "mapping": mapping,
-                    }
-
-        return json_schema
+                schema = {
+                    "oneOf": schemas,
+                    "discriminator": {"propertyName": "kind", "mapping": mapping},
+                }
+            else:
+                schema = handler(core_schema)
+                schema["properties"]["kind"] = {
+                    "const": cls.__name__,
+                    "title": "Kind",
+                    "type": "string",
+                }
+        finally:
+            # Reset temp schema
+            schemas_in_progress.pop(cls)
+        return schema
 
     @classmethod
-    def model_rebuild(
-        cls,
-        *,
-        force=False,
-        raise_errors=True,
-        _parent_namespace_depth=2,
-        _types_namespace=None,
-    ):
-        if cls == DiscriminatedUnionMixin:
-            pass
-        if _is_abstract(cls):
-            subclasses = get_known_concrete_subclasses(cls)
-            kinds = [subclass.__name__ for subclass in subclasses]
-            if kinds:
-                kind_field = cls.model_fields["kind"]
-                kind_field.annotation = Literal[tuple(kinds)]  # type: ignore
-                kind_field.default = kinds[0]
-
-            type_adapter = TypeAdapter(cls.get_serializable_type())
-            cls.__pydantic_core_schema__ = type_adapter.core_schema
-            cls.__pydantic_validator__ = type_adapter.validator
-            cls.__pydantic_serializer__ = type_adapter.serializer
-            return
-
-        return super().model_rebuild(
-            force=force,
-            raise_errors=raise_errors,
-            _parent_namespace_depth=_parent_namespace_depth,
-            _types_namespace=_types_namespace,
+    def resolve_kind(cls, kind: str) -> type[Self]:
+        subclasses = _get_checked_concrete_subclasses(cls)
+        subclass = subclasses.get(kind)
+        if subclass:
+            return subclass
+        raise ValueError(
+            f"Unknown kind '{kind}' for {cls.__module__}.{cls.__name__}; "
+            f"Expected one of: {list(subclasses)}"
         )
 
     @classmethod
@@ -298,7 +279,7 @@ class DiscriminatedUnionMixin(OpenHandsModel, ABC):
         if not _is_abstract(cls):
             return cls
 
-        subclasses = list(get_known_concrete_subclasses(cls))
+        subclasses = _get_checked_concrete_subclasses(cls)
         if not subclasses:
             return cls
 
@@ -306,265 +287,10 @@ class DiscriminatedUnionMixin(OpenHandsModel, ABC):
             # Returning the concrete type ensures Pydantic instantiates the subclass
             # (e.g. Agent) rather than the abstract base (e.g. AgentBase) when there is
             # only ONE concrete subclass.
-            return subclasses[0]
+            return next(iter(subclasses.values()))
 
         serializable_type = Annotated[
-            Union[*tuple(Annotated[t, Tag(t.__name__)] for t in subclasses)],
+            Union[*tuple(Annotated[t, Tag(n)] for n, t in subclasses.items())],
             Discriminator(kind_of),
         ]
         return serializable_type  # type: ignore
-
-    @classmethod
-    def model_validate(cls, obj: Any, **kwargs) -> Self:
-        try:
-            if _is_abstract(cls):
-                resolved = cls.resolve_kind(kind_of(obj))
-            else:
-                resolved = super()
-            result = resolved.model_validate(obj, **kwargs)
-            return result  # type: ignore
-        except ValidationError as e:
-            valid_kinds = [
-                subclass.__name__ for subclass in get_known_concrete_subclasses(cls)
-            ]
-            _handle_discriminated_union_validation_error(e, cls.__name__, valid_kinds)
-
-    @classmethod
-    def model_validate_json(
-        cls,
-        json_data: str | bytes | bytearray,
-        **kwargs,
-    ) -> Self:
-        data = json.loads(json_data)
-        if _is_abstract(cls):
-            resolved = cls.resolve_kind(kind_of(data))
-        else:
-            resolved = super()
-        result = resolved.model_validate(data, **kwargs)
-        return result  # type: ignore
-
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
-
-        # If concrete, stamp kind Literal and collision check
-        if not _is_abstract(cls):
-            # 1) Stamp discriminator
-            cls.kind = cls.__name__
-            cls.__annotations__["kind"] = Literal[cls.__name__]
-
-            # 2) Collision check
-            mro = cls.mro()
-            union_class = mro[mro.index(DiscriminatedUnionMixin) - 1]
-            concretes = get_known_concrete_subclasses(union_class)  # sorted list
-            kinds: dict[str, type] = {}
-            for sub in concretes:
-                k = kind_of(sub)
-                if k in kinds and kinds[k] is not sub:
-                    raise ValueError(
-                        f"Duplicate kind detected for {union_class} : {cls}, {sub}"
-                    )
-                kinds[k] = sub
-
-        # Rebuild any abstract union owners in the MRO that rely on subclass sets
-        for base in cls.mro():
-            # Stop when we pass ourselves
-            if base is cls:
-                continue
-            # Only rebuild abstract DiscriminatedUnion owners
-            if (
-                isinstance(base, type)
-                and issubclass(base, DiscriminatedUnionMixin)
-                and _is_abstract(base)
-            ):
-                base.model_rebuild(force=True)
-
-
-def _rebuild_if_required():
-    if _rebuild_required:
-        rebuild_all()
-
-
-def _extract_discriminated_unions(schema: dict) -> dict:
-    """Extract inline discriminated unions as separate components.
-
-    Recursively scans the schema and extracts any inline discriminated union
-    (oneOf + discriminator + title) as a separate component, replacing it with a $ref.
-    Also deduplicates schemas with identical titles.
-    """
-    import json
-    import re
-    from collections import defaultdict
-
-    if not isinstance(schema, dict):
-        return schema
-
-    # OpenAPI schema names must match this pattern
-    valid_name_pattern = re.compile(r"^[a-zA-Z0-9._-]+$")
-
-    schemas = schema.get("components", {}).get("schemas", {})
-    extracted = {}
-
-    def _find_and_extract(obj, path=""):
-        if not isinstance(obj, dict):
-            return obj
-
-        # Extract inline discriminated unions
-        if "oneOf" in obj and "discriminator" in obj and "title" in obj:
-            title = obj["title"]
-            if (
-                title not in schemas
-                and title not in extracted
-                and valid_name_pattern.match(title)
-            ):
-                extracted[title] = {
-                    "oneOf": obj["oneOf"],
-                    "discriminator": obj["discriminator"],
-                    "title": title,
-                }
-                return {"$ref": f"#/components/schemas/{title}"}
-
-        # Recursively process nested structures
-        result = {}
-        for key, value in obj.items():
-            if isinstance(value, dict):
-                result[key] = _find_and_extract(value, f"{path}.{key}")
-            elif isinstance(value, list):
-                result[key] = [
-                    _find_and_extract(item, f"{path}.{key}[]") for item in value
-                ]
-            else:
-                result[key] = value
-        return result
-
-    schema = _find_and_extract(schema)
-
-    if extracted and "components" in schema and "schemas" in schema["components"]:
-        schema["components"]["schemas"].update(extracted)
-
-    # Deduplicate schemas with same title (prefer *-Output over *-Input over base)
-    schemas = schema.get("components", {}).get("schemas", {})
-    title_to_names = defaultdict(list)
-    for name, defn in schemas.items():
-        if isinstance(defn, dict):
-            title_to_names[defn.get("title", name)].append(name)
-
-    to_remove = {}
-    for title, names in title_to_names.items():
-        if len(names) > 1:
-            # Prefer: *-Output > *-Input > base name
-            keep = sorted(
-                names,
-                key=lambda n: (
-                    0 if n.endswith("-Output") else 1 if n.endswith("-Input") else 2,
-                    n,
-                ),
-            )[0]
-            for name in names:
-                if name != keep:
-                    to_remove[name] = keep
-
-    if to_remove:
-        schema_str = json.dumps(schema)
-        for old, new in to_remove.items():
-            schema_str = schema_str.replace(
-                f'"#/components/schemas/{old}"', f'"#/components/schemas/{new}"'
-            )
-        schema = json.loads(schema_str)
-        for old in to_remove:
-            schema["components"]["schemas"].pop(old, None)
-
-    return schema
-
-
-def _patch_fastapi_discriminated_union_support():
-    """Patch FastAPI to handle discriminated union schemas without $ref.
-
-    This ensures discriminated unions from DiscriminatedUnionMixin work correctly
-    with FastAPI's OpenAPI schema generation. The patch prevents KeyError when
-    FastAPI encounters schemas without $ref keys (which discriminated unions use).
-
-    Also extracts inline discriminated unions as separate schema components for
-    better OpenAPI documentation and Swagger UI display.
-
-    Skips patching if SKIP_FASTAPI_DISCRIMINATED_UNION_FIX environment variable is set.
-    """
-    # Skip patching if environment variable flag is defined
-    if os.environ.get("SKIP_FASTAPI_DISCRIMINATED_UNION_FIX"):
-        logger.debug(
-            "Skipping FastAPI discriminated union patch due to environment variable"
-        )
-        return
-
-    try:
-        import fastapi._compat.v2 as fastapi_v2
-        from fastapi import FastAPI
-
-        _original_remap = fastapi_v2._remap_definitions_and_field_mappings
-
-        def _patched_remap_definitions_and_field_mappings(**kwargs):
-            """Patched version that handles schemas w/o $ref (discriminated unions)."""
-            field_mapping = kwargs.get("field_mapping", {})
-            model_name_map = kwargs.get("model_name_map", {})
-
-            # Build old_name -> new_name map, skipping schemas without $ref
-            old_name_to_new_name_map = {}
-            for field_key, schema in field_mapping.items():
-                model = field_key[0].type_
-                if model not in model_name_map:
-                    continue
-                new_name = model_name_map[model]
-
-                # Skip schemas without $ref (discriminated unions)
-                if "$ref" not in schema:
-                    continue
-
-                old_name = schema["$ref"].split("/")[-1]
-                if old_name in {f"{new_name}-Input", f"{new_name}-Output"}:
-                    continue
-                old_name_to_new_name_map[old_name] = new_name
-
-            # Replace refs using FastAPI's helper
-            from fastapi._compat.v2 import _replace_refs
-
-            new_field_mapping = {}
-            for field_key, schema in field_mapping.items():
-                new_schema = _replace_refs(
-                    schema=schema,
-                    old_name_to_new_name_map=old_name_to_new_name_map,
-                )
-                new_field_mapping[field_key] = new_schema
-
-            definitions = kwargs.get("definitions", {})
-            new_definitions = {}
-            for key, value in definitions.items():
-                new_key = old_name_to_new_name_map.get(key, key)
-                new_value = _replace_refs(
-                    schema=value,
-                    old_name_to_new_name_map=old_name_to_new_name_map,
-                )
-                new_definitions[new_key] = new_value
-
-            return new_field_mapping, new_definitions
-
-        # Apply the patch
-        fastapi_v2._remap_definitions_and_field_mappings = (
-            _patched_remap_definitions_and_field_mappings
-        )
-
-        # Patch FastAPI.openapi() to extract discriminated unions
-        _original_openapi = FastAPI.openapi
-
-        def _patched_openapi(self):
-            """Patched openapi() that extracts discriminated unions."""
-            schema = _original_openapi(self)
-            return _extract_discriminated_unions(schema)
-
-        FastAPI.openapi = _patched_openapi
-
-    except (ImportError, AttributeError):
-        # FastAPI not available or internal API changed
-        pass
-
-
-# Always call the FastAPI patch after DiscriminatedUnionMixin definition
-_patch_fastapi_discriminated_union_support()