Merge pull request #448 from jlowin/openapi-def

Ensure openapi defs for structured objects are loaded properly

src/fastmcp/prompts/prompt.py

@@ -13,7 +13,7 @@ from mcp.types import PromptArgument as MCPPromptArgument
 from pydantic import BaseModel, BeforeValidator, Field, TypeAdapter, validate_call
 
 from fastmcp.server.dependencies import get_context
-from fastmcp.utilities.json_schema import prune_params
+from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.types import (
     _convert_set_defaults,
@@ -115,7 +115,11 @@ class Prompt(BaseModel):
 
         context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
         if context_kwarg:
-            parameters = prune_params(parameters, params=[context_kwarg])
+            prune_params = [context_kwarg]
+        else:
+            prune_params = None
+
+        parameters = compress_schema(parameters, prune_params=prune_params)
 
         # Convert parameters to PromptArguments
         arguments: list[PromptArgument] = []

src/fastmcp/resources/template.py

@@ -21,6 +21,7 @@ from pydantic import (
 
 from fastmcp.resources.types import FunctionResource, Resource
 from fastmcp.server.dependencies import get_context
+from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.types import (
     _convert_set_defaults,
     find_kwarg_by_type,
@@ -150,6 +151,10 @@ class ResourceTemplate(BaseModel):
         # Get schema from TypeAdapter - will fail if function isn't properly typed
         parameters = TypeAdapter(fn).json_schema()
 
+        # compress the schema
+        prune_params = [context_kwarg] if context_kwarg else None
+        parameters = compress_schema(parameters, prune_params=prune_params)
+
         # ensure the arguments are properly cast
         fn = validate_call(fn)
 

src/fastmcp/tools/tool.py

@@ -12,7 +12,7 @@ from pydantic import BaseModel, BeforeValidator, Field
 
 import fastmcp
 from fastmcp.server.dependencies import get_context
-from fastmcp.utilities.json_schema import prune_params
+from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.types import (
     Image,
@@ -81,7 +81,11 @@ class Tool(BaseModel):
 
         context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
         if context_kwarg:
-            schema = prune_params(schema, params=[context_kwarg])
+            prune_params = [context_kwarg]
+        else:
+            prune_params = None
+
+        schema = compress_schema(schema, prune_params=prune_params)
 
         return cls(
             fn=fn,

src/fastmcp/utilities/json_schema.py

@@ -14,6 +14,7 @@ def _prune_param(schema: dict, param: str) -> dict:
     removed = props.pop(param, None)
     if removed is None:  # nothing to do
         return schema
+
     # Keep empty properties object rather than removing it entirely
     schema["properties"] = props
     if param in schema.get("required", []):
@@ -21,7 +22,12 @@ def _prune_param(schema: dict, param: str) -> dict:
         if not schema["required"]:
             schema.pop("required")
 
-    # ── 2. collect all remaining local $ref targets ───────────────────
+    return schema
+
+
+def _prune_unused_defs(schema: dict) -> dict:
+    """Remove unused definitions from the schema."""
+    # collect all remaining local $ref targets
     used_defs: set[str] = set()
 
     def walk(node: object) -> None:  # depth-first traversal
@@ -37,7 +43,8 @@ def _prune_param(schema: dict, param: str) -> dict:
 
     walk(schema)
 
-    # ── 3. remove orphaned definitions ────────────────────────────────
+    # remove orphaned definitions
+
     defs = schema.get("$defs", {})
     for def_name in list(defs):
         if def_name not in used_defs:
@@ -48,12 +55,28 @@ def _prune_param(schema: dict, param: str) -> dict:
     return schema
 
 
-def prune_params(schema: dict, params: list[str]) -> dict:
+def _prune_additional_properties(schema: dict) -> dict:
+    """Remove additionalProperties from the schema if it is False."""
+    if schema.get("additionalProperties", None) is False:
+        schema.pop("additionalProperties")
+    return schema
+
+
+def compress_schema(
+    schema: dict,
+    prune_params: list[str] | None = None,
+    prune_defs: bool = True,
+    prune_additional_properties: bool = True,
+) -> dict:
     """
     Remove the given parameters from the schema.
 
     """
     schema = copy.deepcopy(schema)
-    for param in params:
+    for param in prune_params or []:
         schema = _prune_param(schema, param=param)
+    if prune_defs:
+        schema = _prune_unused_defs(schema)
+    if prune_additional_properties:
+        schema = _prune_additional_properties(schema)
     return schema

src/fastmcp/utilities/openapi.py

@@ -84,6 +84,9 @@ class HTTPRoute(BaseModel):
     responses: dict[str, ResponseInfo] = Field(
         default_factory=dict
     )  # Key: status code str
+    schema_definitions: dict[str, JsonSchema] = Field(
+        default_factory=dict
+    )  # Store component schemas
 
 
 # Export public symbols
@@ -221,6 +224,27 @@ class OpenAPI31Parser(BaseOpenAPIParser):
             logger.warning("OpenAPI schema has no paths defined.")
             return []
 
+        # Extract component schemas to add to each route
+        schema_definitions = {}
+        if hasattr(self.openapi, "components") and self.openapi.components:
+            components = self.openapi.components
+            if hasattr(components, "schemas") and components.schemas:
+                for name, schema in components.schemas.items():
+                    try:
+                        if isinstance(schema, Reference):
+                            resolved_schema = self._resolve_ref(schema)
+                            schema_definitions[name] = self._extract_schema_as_dict(
+                                resolved_schema
+                            )
+                        else:
+                            schema_definitions[name] = self._extract_schema_as_dict(
+                                schema
+                            )
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to extract schema definition '{name}': {e}"
+                        )
+
         for path_str, path_item_obj in self.openapi.paths.items():
             if not isinstance(path_item_obj, PathItem):
                 logger.warning(
@@ -269,6 +293,7 @@ class OpenAPI31Parser(BaseOpenAPIParser):
                             parameters=parameters,
                             request_body=request_body_info,
                             responses=responses,
+                            schema_definitions=schema_definitions,
                         )
                         routes.append(route)
                         logger.info(
@@ -386,16 +411,36 @@ class OpenAPI31Parser(BaseOpenAPIParser):
 
                 param_schema_dict = {}
                 if param_schema_obj:  # Check if schema exists
+                    # Resolve the schema if it's a reference
+                    resolved_schema = self._resolve_ref(param_schema_obj)
                     param_schema_dict = self._extract_schema_as_dict(param_schema_obj)
+
+                    # Ensure default value is preserved from resolved schema
+                    if (
+                        not isinstance(resolved_schema, Reference)
+                        and hasattr(resolved_schema, "default")
+                        and resolved_schema.default is not None
+                    ):
+                        param_schema_dict["default"] = resolved_schema.default
                 elif parameter.content:
                     # Handle complex parameters with 'content'
                     first_media_type = next(iter(parameter.content.values()), None)
                     if (
                         first_media_type and first_media_type.media_type_schema
                     ):  # CORRECTED: Use 'media_type_schema'
-                        param_schema_dict = self._extract_schema_as_dict(
-                            first_media_type.media_type_schema
-                        )
+                        # Resolve the schema if it's a reference
+                        media_schema = first_media_type.media_type_schema
+                        resolved_media_schema = self._resolve_ref(media_schema)
+                        param_schema_dict = self._extract_schema_as_dict(media_schema)
+
+                        # Ensure default value is preserved from resolved schema
+                        if (
+                            not isinstance(resolved_media_schema, Reference)
+                            and hasattr(resolved_media_schema, "default")
+                            and resolved_media_schema.default is not None
+                        ):
+                            param_schema_dict["default"] = resolved_media_schema.default
+
                         logger.debug(
                             f"Parameter '{parameter.name}' using schema from 'content' field."
                         )
@@ -543,6 +588,27 @@ class OpenAPI30Parser(BaseOpenAPIParser):
             logger.warning("OpenAPI schema has no paths defined.")
             return []
 
+        # Extract component schemas to add to each route
+        schema_definitions = {}
+        if hasattr(self.openapi, "components") and self.openapi.components:
+            components = self.openapi.components
+            if hasattr(components, "schemas") and components.schemas:
+                for name, schema in components.schemas.items():
+                    try:
+                        if isinstance(schema, Reference_30):
+                            resolved_schema = self._resolve_ref(schema)
+                            schema_definitions[name] = self._extract_schema_as_dict(
+                                resolved_schema
+                            )
+                        else:
+                            schema_definitions[name] = self._extract_schema_as_dict(
+                                schema
+                            )
+                    except Exception as e:
+                        logger.warning(
+                            f"Failed to extract schema definition '{name}': {e}"
+                        )
+
         for path_str, path_item_obj in self.openapi.paths.items():
             if not isinstance(path_item_obj, PathItem_30):
                 logger.warning(
@@ -593,6 +659,7 @@ class OpenAPI30Parser(BaseOpenAPIParser):
                             parameters=parameters,
                             request_body=request_body_info,
                             responses=responses,
+                            schema_definitions=schema_definitions,
                         )
                         routes.append(route)
                         logger.info(
@@ -711,14 +778,34 @@ class OpenAPI30Parser(BaseOpenAPIParser):
 
                 param_schema_dict = {}
                 if param_schema_obj:  # Check if schema exists
+                    # Resolve the schema if it's a reference
+                    resolved_schema = self._resolve_ref(param_schema_obj)
                     param_schema_dict = self._extract_schema_as_dict(param_schema_obj)
+
+                    # Ensure default value is preserved from resolved schema
+                    if (
+                        not isinstance(resolved_schema, Reference_30)
+                        and hasattr(resolved_schema, "default")
+                        and resolved_schema.default is not None
+                    ):
+                        param_schema_dict["default"] = resolved_schema.default
                 elif parameter.content:
                     # Handle complex parameters with 'content'
                     first_media_type = next(iter(parameter.content.values()), None)
                     if first_media_type and first_media_type.media_type_schema:
-                        param_schema_dict = self._extract_schema_as_dict(
-                            first_media_type.media_type_schema
-                        )
+                        # Resolve the schema if it's a reference
+                        media_schema = first_media_type.media_type_schema
+                        resolved_media_schema = self._resolve_ref(media_schema)
+                        param_schema_dict = self._extract_schema_as_dict(media_schema)
+
+                        # Ensure default value is preserved from resolved schema
+                        if (
+                            not isinstance(resolved_media_schema, Reference_30)
+                            and hasattr(resolved_media_schema, "default")
+                            and resolved_media_schema.default is not None
+                        ):
+                            param_schema_dict["default"] = resolved_media_schema.default
+
                         logger.debug(
                             f"Parameter '{parameter.name}' using schema from 'content' field."
                         )
@@ -1173,6 +1260,23 @@ def _combine_schemas(route: openapi.HTTPRoute) -> dict[str, Any]:
         # 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
@@ -1193,8 +1297,19 @@ def _combine_schemas(route: openapi.HTTPRoute) -> dict[str, Any]:
         if route.request_body.required:
             required.extend(body_schema.get("required", []))
 
-    return {
+    result = {
         "type": "object",
         "properties": properties,
         "required": required,
     }
+
+    # Add schema definitions if available
+    if route.schema_definitions:
+        result["$defs"] = route.schema_definitions
+
+    # Use compress_schema to remove unused definitions
+    from fastmcp.utilities.json_schema import compress_schema
+
+    result = compress_schema(result)
+
+    return result

tests/server/test_openapi.py

@@ -1,6 +1,7 @@
 import base64
 import json
 import re
+from enum import Enum
 
 import httpx
 import pytest
@@ -1817,3 +1818,75 @@ class TestReprMethods:
         assert f"name={template.name!r}" in template_repr
         assert "uri_template=" in template_repr
         assert "path=" in template_repr
+
+
+class TestEnumHandling:
+    """Tests for handling enum parameters in OpenAPI schemas."""
+
+    async def test_enum_parameter_schema(self):
+        """Test that enum parameters are properly handled in tool parameter schemas."""
+
+        # Define an enum just like in example.py
+        class QueryEnum(str, Enum):
+            foo = "foo"
+            bar = "bar"
+            baz = "baz"
+
+        # Create a minimal FastAPI app with an endpoint using the enum
+        app = FastAPI()
+
+        @app.post("/items/{item_id}")
+        def read_item(
+            item_id: int,
+            query: QueryEnum | None = None,
+        ):
+            return {"item_id": item_id, "query": query}
+
+        # Create a client for the app
+        client = AsyncClient(transport=ASGITransport(app=app), base_url="http://test")
+
+        # Create the FastMCPOpenAPI server from the app
+        openapi_spec = app.openapi()
+        server = FastMCPOpenAPI(
+            openapi_spec=openapi_spec,
+            client=client,
+            name="Enum Test",
+        )
+
+        # Get the tools from the server
+        tools = server._tool_manager.list_tools()
+
+        # Find the read_item tool
+        read_item_tool = next(
+            (t for t in tools if t.name == "read_item_items__item_id__post"), None
+        )
+
+        # Verify the tool exists
+        assert read_item_tool is not None, "read_item tool wasn't created"
+
+        # Check that the parameters include the enum reference
+        assert "properties" in read_item_tool.parameters
+        assert "query" in read_item_tool.parameters["properties"]
+
+        # Check for the anyOf with $ref to the enum definition
+        query_param = read_item_tool.parameters["properties"]["query"]
+        assert "anyOf" in query_param
+
+        # Find the ref in the anyOf list
+        ref_found = False
+        for option in query_param["anyOf"]:
+            if "$ref" in option and option["$ref"].startswith("#/$defs/QueryEnum"):
+                ref_found = True
+                break
+
+        assert ref_found, "Reference to enum definition not found in query parameter"
+
+        # Check that the $defs section exists and contains the enum definition
+        assert "$defs" in read_item_tool.parameters
+        assert "QueryEnum" in read_item_tool.parameters["$defs"]
+
+        # Verify the enum definition
+        enum_def = read_item_tool.parameters["$defs"]["QueryEnum"]
+        assert "enum" in enum_def
+        assert enum_def["enum"] == ["foo", "bar", "baz"]
+        assert enum_def["type"] == "string"

tests/utilities/test_json_schema.py

@@ -1,110 +1,246 @@
-from fastmcp.utilities.json_schema import _prune_param, prune_params
-
-
-def test_prune_param_nonexistent():
-    """Test pruning a parameter that doesn't exist."""
-    schema = {"properties": {"foo": {"type": "string"}}}
-    result = _prune_param(schema, "bar")
-    assert result == schema  # Schema should be unchanged
-
-
-def test_prune_param_exists():
-    """Test pruning a parameter that exists."""
-    schema = {"properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}}
-    result = _prune_param(schema, "bar")
-    assert result["properties"] == {"foo": {"type": "string"}}
-
-
-def test_prune_param_last_property():
-    """Test pruning the only/last parameter, should leave empty properties object."""
-    schema = {"properties": {"foo": {"type": "string"}}}
-    result = _prune_param(schema, "foo")
-    assert "properties" in result
-    assert result["properties"] == {}
-
-
-def test_prune_param_from_required():
-    """Test pruning a parameter that's in the required list."""
-    schema = {
-        "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}},
-        "required": ["foo", "bar"],
-    }
-    result = _prune_param(schema, "bar")
-    assert result["required"] == ["foo"]
-
-
-def test_prune_param_last_required():
-    """Test pruning the last required parameter, should remove required field."""
-    schema = {
-        "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}},
-        "required": ["foo"],
-    }
-    result = _prune_param(schema, "foo")
-    assert "required" not in result
-
-
-def test_prune_param_with_refs():
-    """Test pruning a parameter that has references in $defs."""
-    schema = {
-        "properties": {
-            "foo": {"$ref": "#/$defs/foo_def"},
-            "bar": {"$ref": "#/$defs/bar_def"},
-        },
-        "$defs": {
-            "foo_def": {"type": "string"},
-            "bar_def": {"type": "integer"},
-        },
-    }
-    result = _prune_param(schema, "bar")
-    assert "bar_def" not in result["$defs"]
-    assert "foo_def" in result["$defs"]
-
-
-def test_prune_param_all_refs():
-    """Test pruning all parameters with refs, should remove $defs."""
-    schema = {
-        "properties": {
-            "foo": {"$ref": "#/$defs/foo_def"},
-        },
-        "$defs": {
-            "foo_def": {"type": "string"},
-        },
-    }
-    result = _prune_param(schema, "foo")
-    assert "$defs" not in result
-
-
-def test_prune_params_multiple():
-    """Test pruning multiple parameters at once."""
-    schema = {
-        "properties": {
-            "foo": {"type": "string"},
-            "bar": {"type": "integer"},
-            "baz": {"type": "boolean"},
-        },
-        "required": ["foo", "bar"],
-    }
-    result = prune_params(schema, ["foo", "baz"])
-    assert result["properties"] == {"bar": {"type": "integer"}}
-    assert result["required"] == ["bar"]
-
-
-def test_prune_params_nested_refs():
-    """Test pruning with nested references."""
-    schema = {
-        "properties": {
-            "foo": {
-                "type": "object",
-                "properties": {"nested": {"$ref": "#/$defs/nested_def"}},
+from fastmcp.utilities.json_schema import (
+    _prune_additional_properties,
+    _prune_param,
+    _prune_unused_defs,
+    compress_schema,
+)
+
+
+class TestPruneParam:
+    """Tests for the _prune_param function."""
+
+    def test_nonexistent(self):
+        """Test pruning a parameter that doesn't exist."""
+        schema = {"properties": {"foo": {"type": "string"}}}
+        result = _prune_param(schema, "bar")
+        assert result == schema  # Schema should be unchanged
+
+    def test_exists(self):
+        """Test pruning a parameter that exists."""
+        schema = {"properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}}}
+        result = _prune_param(schema, "bar")
+        assert result["properties"] == {"foo": {"type": "string"}}
+
+    def test_last_property(self):
+        """Test pruning the only/last parameter, should leave empty properties object."""
+        schema = {"properties": {"foo": {"type": "string"}}}
+        result = _prune_param(schema, "foo")
+        assert "properties" in result
+        assert result["properties"] == {}
+
+    def test_from_required(self):
+        """Test pruning a parameter that's in the required list."""
+        schema = {
+            "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}},
+            "required": ["foo", "bar"],
+        }
+        result = _prune_param(schema, "bar")
+        assert result["required"] == ["foo"]
+
+    def test_last_required(self):
+        """Test pruning the last required parameter, should remove required field."""
+        schema = {
+            "properties": {"foo": {"type": "string"}, "bar": {"type": "integer"}},
+            "required": ["foo"],
+        }
+        result = _prune_param(schema, "foo")
+        assert "required" not in result
+
+
+class TestPruneUnusedDefs:
+    """Tests for the _prune_unused_defs function."""
+
+    def test_removes_unreferenced_defs(self):
+        """Test that unreferenced definitions are removed."""
+        schema = {
+            "properties": {
+                "foo": {"$ref": "#/$defs/foo_def"},
             },
-            "bar": {"$ref": "#/$defs/bar_def"},
-        },
-        "$defs": {
-            "nested_def": {"type": "string"},
-            "bar_def": {"type": "integer"},
-        },
-    }
-    # Removing foo should keep nested_def as it's not referenced anymore
-    result = _prune_param(schema, "foo")
-    assert "nested_def" not in result["$defs"]
-    assert "bar_def" in result["$defs"]
+            "$defs": {
+                "foo_def": {"type": "string"},
+                "unused_def": {"type": "integer"},
+            },
+        }
+        result = _prune_unused_defs(schema)
+        assert "foo_def" in result["$defs"]
+        assert "unused_def" not in result["$defs"]
+
+    def test_nested_references_kept(self):
+        """Test that definitions referenced via nesting are kept."""
+        schema = {
+            "properties": {
+                "foo": {"$ref": "#/$defs/foo_def"},
+            },
+            "$defs": {
+                "foo_def": {
+                    "type": "object",
+                    "properties": {"nested": {"$ref": "#/$defs/nested_def"}},
+                },
+                "nested_def": {"type": "string"},
+                "unused_def": {"type": "integer"},
+            },
+        }
+        result = _prune_unused_defs(schema)
+        assert "foo_def" in result["$defs"]
+        assert "nested_def" in result["$defs"]
+        assert "unused_def" not in result["$defs"]
+
+    def test_array_references_kept(self):
+        """Test that definitions referenced in array items are kept."""
+        schema = {
+            "properties": {
+                "items": {"type": "array", "items": {"$ref": "#/$defs/item_def"}},
+            },
+            "$defs": {
+                "item_def": {"type": "string"},
+                "unused_def": {"type": "integer"},
+            },
+        }
+        result = _prune_unused_defs(schema)
+        assert "item_def" in result["$defs"]
+        assert "unused_def" not in result["$defs"]
+
+    def test_removes_defs_field_when_empty(self):
+        """Test that $defs field is removed when all definitions are unused."""
+        schema = {
+            "properties": {
+                "foo": {"type": "string"},
+            },
+            "$defs": {
+                "unused_def": {"type": "integer"},
+            },
+        }
+        result = _prune_unused_defs(schema)
+        assert "$defs" not in result
+
+
+class TestPruneAdditionalProperties:
+    """Tests for the _prune_additional_properties function."""
+
+    def test_removes_when_false(self):
+        """Test that additionalProperties is removed when it's false."""
+        schema = {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}},
+            "additionalProperties": False,
+        }
+        result = _prune_additional_properties(schema)
+        assert "additionalProperties" not in result
+
+    def test_keeps_when_true(self):
+        """Test that additionalProperties is kept when it's true."""
+        schema = {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}},
+            "additionalProperties": True,
+        }
+        result = _prune_additional_properties(schema)
+        assert "additionalProperties" in result
+        assert result["additionalProperties"] is True
+
+    def test_keeps_when_object(self):
+        """Test that additionalProperties is kept when it's an object schema."""
+        schema = {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}},
+            "additionalProperties": {"type": "string"},
+        }
+        result = _prune_additional_properties(schema)
+        assert "additionalProperties" in result
+        assert result["additionalProperties"] == {"type": "string"}
+
+
+class TestCompressSchema:
+    """Tests for the compress_schema function."""
+
+    def test_prune_params(self):
+        """Test pruning parameters with compress_schema."""
+        schema = {
+            "properties": {
+                "foo": {"type": "string"},
+                "bar": {"type": "integer"},
+                "baz": {"type": "boolean"},
+            },
+            "required": ["foo", "bar"],
+        }
+        result = compress_schema(schema, prune_params=["foo", "baz"])
+        assert result["properties"] == {"bar": {"type": "integer"}}
+        assert result["required"] == ["bar"]
+
+    def test_prune_defs(self):
+        """Test pruning unused definitions with compress_schema."""
+        schema = {
+            "properties": {
+                "foo": {"$ref": "#/$defs/foo_def"},
+                "bar": {"type": "integer"},
+            },
+            "$defs": {
+                "foo_def": {"type": "string"},
+                "unused_def": {"type": "number"},
+            },
+        }
+        result = compress_schema(schema)
+        assert "foo_def" in result["$defs"]
+        assert "unused_def" not in result["$defs"]
+
+    def test_disable_prune_defs(self):
+        """Test disabling pruning of unused definitions."""
+        schema = {
+            "properties": {
+                "foo": {"$ref": "#/$defs/foo_def"},
+                "bar": {"type": "integer"},
+            },
+            "$defs": {
+                "foo_def": {"type": "string"},
+                "unused_def": {"type": "number"},
+            },
+        }
+        result = compress_schema(schema, prune_defs=False)
+        assert "foo_def" in result["$defs"]
+        assert "unused_def" in result["$defs"]
+
+    def test_pruning_additional_properties(self):
+        """Test pruning additionalProperties when False."""
+        schema = {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}},
+            "additionalProperties": False,
+        }
+        result = compress_schema(schema)
+        assert "additionalProperties" not in result
+
+    def test_disable_pruning_additional_properties(self):
+        """Test disabling pruning of additionalProperties."""
+        schema = {
+            "type": "object",
+            "properties": {"foo": {"type": "string"}},
+            "additionalProperties": False,
+        }
+        result = compress_schema(schema, prune_additional_properties=False)
+        assert "additionalProperties" in result
+        assert result["additionalProperties"] is False
+
+    def test_combined_operations(self):
+        """Test all pruning operations together."""
+        schema = {
+            "type": "object",
+            "properties": {
+                "keep": {"type": "string"},
+                "remove": {"$ref": "#/$defs/remove_def"},
+            },
+            "required": ["keep", "remove"],
+            "additionalProperties": False,
+            "$defs": {
+                "remove_def": {"type": "string"},
+                "unused_def": {"type": "number"},
+            },
+        }
+        result = compress_schema(schema, prune_params=["remove"])
+        # Check that parameter was removed
+        assert "remove" not in result["properties"]
+        # Check that required list was updated
+        assert result["required"] == ["keep"]
+        # Check that unused definitions were removed
+        assert "$defs" not in result  # Both defs should be gone
+        # Check that additionalProperties was removed
+        assert "additionalProperties" not in result

tests/utilities/test_typeadapter.py

@@ -13,7 +13,7 @@ import annotated_types
 import pytest
 from pydantic import BaseModel, Field
 
-from fastmcp.utilities.json_schema import prune_params
+from fastmcp.utilities.json_schema import compress_schema
 from fastmcp.utilities.types import get_cached_typeadapter
 
 
@@ -175,7 +175,7 @@ def test_skip_names():
     # Get schema and prune parameters
     type_adapter = get_cached_typeadapter(func_with_many_params)
     schema = type_adapter.json_schema()
-    pruned_schema = prune_params(schema, params=["skip_this", "also_skip"])
+    pruned_schema = compress_schema(schema, prune_params=["skip_this", "also_skip"])
 
     # Check that only the desired parameters remain
     assert "keep_this" in pruned_schema["properties"]