Merge pull request #697 from phateffect/replace-w-defs

replace $ref pointing to `#/components/schemas/` with `#/$defs/`

src/fastmcp/utilities/openapi.py

@@ -872,6 +872,50 @@ def format_description_with_responses(
     return "\n".join(desc_parts)
 
 
+def _replace_ref_with_defs(
+    info: dict[str, Any], description: str | None = None
+) -> dict[str, Any]:
+    """
+    Replace openapi $ref with jsonschema $defs
+
+    Examples:
+    - {"type": "object", "properties": {"$ref": "#/components/schemas/..."}}
+    - {"$ref": "#/components/schemas/..."}
+    - {"items": {"$ref": "#/components/schemas/..."}}
+    - {"anyOf": [{"$ref": "#/components/schemas/..."}]}
+    - {"allOf": [{"$ref": "#/components/schemas/..."}]}
+    - {"oneOf": [{"$ref": "#/components/schemas/..."}]}
+
+    Args:
+        info: dict[str, Any]
+        description: str | None
+
+    Returns:
+        dict[str, Any]
+    """
+    schema = info.copy()
+    if ref_path := schema.get("$ref"):
+        if ref_path.startswith("#/components/schemas/"):
+            schema_name = ref_path.split("/")[-1]
+            schema["$ref"] = f"#/$defs/{schema_name}"
+    elif properties := schema.get("properties"):
+        if "$ref" in properties:
+            schema["properties"] = _replace_ref_with_defs(properties)
+        else:
+            schema["properties"] = {
+                prop_name: _replace_ref_with_defs(prop_schema)
+                for prop_name, prop_schema in properties.items()
+            }
+    elif item_schema := schema.get("items"):
+        schema["items"] = _replace_ref_with_defs(item_schema)
+    for section in ["anyOf", "allOf", "oneOf"]:
+        for i, item in enumerate(schema.get(section, [])):
+            schema[section][i] = _replace_ref_with_defs(item)
+    if info.get("description", description) and not schema.get("description"):
+        schema["description"] = description
+    return schema
+
+
 def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
     """
     Combines parameter and request body schemas into a single schema.
@@ -889,38 +933,18 @@ def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
     for param in route.parameters:
         if param.required:
             required.append(param.name)
-
-        # Copy the schema and add description if available
-        param_schema = param.schema_.copy() if isinstance(param.schema_, dict) else {}
-
-        # Convert #/components/schemas references to #/$defs references
-        if isinstance(param_schema, dict) and "$ref" in param_schema:
-            ref_path = param_schema["$ref"]
-            if ref_path.startswith("#/components/schemas/"):
-                schema_name = ref_path.split("/")[-1]
-                param_schema["$ref"] = f"#/$defs/{schema_name}"
-
-        # Also handle anyOf, allOf, oneOf references
-        for section in ["anyOf", "allOf", "oneOf"]:
-            if section in param_schema and isinstance(param_schema[section], list):
-                for i, item in enumerate(param_schema[section]):
-                    if isinstance(item, dict) and "$ref" in item:
-                        ref_path = item["$ref"]
-                        if ref_path.startswith("#/components/schemas/"):
-                            schema_name = ref_path.split("/")[-1]
-                            param_schema[section][i]["$ref"] = f"#/$defs/{schema_name}"
-
-        # Add parameter description to schema if available and not already present
-        if param.description and not param_schema.get("description"):
-            param_schema["description"] = param.description
-
-        properties[param.name] = param_schema
+        properties[param.name] = _replace_ref_with_defs(
+            param.schema_.copy(), param.description
+        )
 
     # Add request body if it exists
     if route.request_body and route.request_body.content_schema:
         # For now, just use the first content type's schema
         content_type = next(iter(route.request_body.content_schema))
-        body_schema = route.request_body.content_schema[content_type]
+        body_schema = _replace_ref_with_defs(
+            route.request_body.content_schema[content_type].copy(),
+            route.request_body.description,
+        )
         body_props = body_schema.get("properties", {})
 
         # Add request body properties
@@ -935,7 +959,6 @@ def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
         "properties": properties,
         "required": required,
     }
-
     # Add schema definitions if available
     if route.schema_definitions:
         result["$defs"] = route.schema_definitions

tests/utilities/openapi/test_openapi.py

@@ -6,7 +6,11 @@ import pytest
 from fastapi import Body, FastAPI, Path, Query
 from pydantic import BaseModel, Field
 
-from fastmcp.utilities.openapi import parse_openapi_to_http_routes
+from fastmcp.utilities.openapi import (
+    _combine_schemas,
+    _replace_ref_with_defs,
+    parse_openapi_to_http_routes,
+)
 
 # --- Test Data: Static OpenAPI Schema Dictionaries --- #
 
@@ -1023,6 +1027,9 @@ def test_openapi_30_reference_resolution(openapi_30_with_references):
     # or it still has a $ref field
     assert "properties" in category or "$ref" in category
 
+    combined_schema = _combine_schemas(route)
+    assert "#/$defs/" in combined_schema["properties"]["category"]["$ref"]
+
 
 def test_openapi_31_reference_resolution(openapi_31_with_references):
     """Test that references are correctly resolved in OpenAPI 3.1 schemas."""
@@ -1057,6 +1064,9 @@ def test_openapi_31_reference_resolution(openapi_31_with_references):
     # or it still has a $ref field
     assert "properties" in category or "$ref" in category
 
+    combined_schema = _combine_schemas(route)
+    assert "#/$defs/" in combined_schema["properties"]["category"]["$ref"]
+
 
 def test_consistent_output_across_versions(
     openapi_30_with_references, openapi_31_with_references
@@ -1093,3 +1103,103 @@ def test_consistent_output_across_versions(
         "properties"
     ]
     assert set(schema_30.keys()) == set(schema_31.keys())
+
+
+class TestReplaceRefWithDefs:
+    @pytest.fixture(scope="class")
+    def schemas(self):
+        """Provide test schemas for _replace_ref_with_defs function."""
+        return {
+            "ref_type": {
+                "$ref": "#/components/schemas/RefFoo",
+            },
+            "object_type": {
+                "type": "object",
+                "properties": {"$ref": "#/components/schemas/ObjectFoo"},
+            },
+            "array_type": {
+                "type": "array",
+                "items": {"$ref": "#/components/schemas/ArrayFoo"},
+            },
+            "any_of_type": {
+                "anyOf": [
+                    {"$ref": "#/components/schemas/AnyOfFoo"},
+                    {"$ref": "#/components/schemas/AnyOfBar"},
+                ]
+            },
+            "all_of_type": {
+                "allOf": [
+                    {"$ref": "#/components/schemas/AllOfFoo"},
+                    {"$ref": "#/components/schemas/AllOfBar"},
+                ]
+            },
+            "one_of_type": {
+                "oneOf": [
+                    {"$ref": "#/components/schemas/OneOfFoo"},
+                    {"$ref": "#/components/schemas/OneOfBar"},
+                ]
+            },
+            "nested_type": {
+                "type": "object",
+                "properties": {
+                    "pets": {
+                        "oneOf": [
+                            {"$ref": "#/components/schemas/Cat"},
+                            {"$ref": "#/components/schemas/Dog"},
+                        ]
+                    },
+                },
+            },
+        }
+
+    def test_replace_direct_ref(self, schemas):
+        """Test replacing direct $ref references."""
+        result = _replace_ref_with_defs(schemas["ref_type"])
+        assert result == {"$ref": "#/$defs/RefFoo"}
+
+    def test_replace_object_property_ref(self, schemas):
+        """Test replacing $ref in object properties."""
+        result = _replace_ref_with_defs(schemas["object_type"])
+        assert result == {
+            "type": "object",
+            "properties": {"$ref": "#/$defs/ObjectFoo"},
+        }
+
+    def test_replace_array_items_ref(self, schemas):
+        """Test replacing $ref in array items."""
+        result = _replace_ref_with_defs(schemas["array_type"])
+        assert result == {
+            "type": "array",
+            "items": {"$ref": "#/$defs/ArrayFoo"},
+        }
+
+    def test_replace_any_of_refs(self, schemas):
+        """Test replacing $ref in anyOf schemas."""
+        result = _replace_ref_with_defs(schemas["any_of_type"])
+        assert result == {
+            "anyOf": [{"$ref": "#/$defs/AnyOfFoo"}, {"$ref": "#/$defs/AnyOfBar"}]
+        }
+
+    def test_replace_all_of_refs(self, schemas):
+        """Test replacing $ref in allOf schemas."""
+        result = _replace_ref_with_defs(schemas["all_of_type"])
+        assert result == {
+            "allOf": [{"$ref": "#/$defs/AllOfFoo"}, {"$ref": "#/$defs/AllOfBar"}]
+        }
+
+    def test_replace_one_of_refs(self, schemas):
+        """Test replacing $ref in oneOf schemas."""
+        result = _replace_ref_with_defs(schemas["one_of_type"])
+        assert result == {
+            "oneOf": [{"$ref": "#/$defs/OneOfFoo"}, {"$ref": "#/$defs/OneOfBar"}]
+        }
+
+    def test_replace_nested_refs(self, schemas):
+        """Test replacing $ref in deeply nested schema structures."""
+        result = _replace_ref_with_defs(schemas["nested_type"])
+        assert result == {
+            "type": "object",
+            "properties": {
+                "pets": {"oneOf": [{"$ref": "#/$defs/Cat"}, {"$ref": "#/$defs/Dog"}]}
+            },
+        }