diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py index cab7adba0..61ecf6907 100644 --- a/src/fastmcp/tools/tool_transform.py +++ b/src/fastmcp/tools/tool_transform.py @@ -13,6 +13,7 @@ from pydantic.functional_validators import BeforeValidator from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult, _convert_to_content from fastmcp.utilities.components import _convert_set_default_none +from fastmcp.utilities.json_schema import compress_schema from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import ( FastMCPBaseModel, @@ -645,6 +646,7 @@ class TransformedTool(Tool): if parent_defs: schema["$defs"] = parent_defs + schema = compress_schema(schema, prune_defs=True) # Create forwarding function that closes over everything it needs async def _forward(**kwargs): diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py index e65154ec8..f498d01ad 100644 --- a/tests/tools/test_tool_transform.py +++ b/tests/tools/test_tool_transform.py @@ -192,6 +192,31 @@ async def test_hide_required_param_with_user_default_works(): assert result.structured_content == {"result": 25} +async def test_hidden_param_prunes_defs(): + class VisibleType(BaseModel): + x: int + + class HiddenType(BaseModel): + y: int + + @Tool.from_function + def tool_with_refs(a: VisibleType, b: HiddenType | None = None) -> int: + return a.x + (b.y if b else 0) + + # Hide parameter 'b' + new_tool = Tool.from_tool( + tool_with_refs, transform_args={"b": ArgTransform(hide=True)} + ) + + schema = new_tool.parameters + # Only 'a' should be visible + assert list(schema["properties"].keys()) == ["a"] + # $defs should only contain VisibleType, not HiddenType + defs = schema.get("$defs", {}) + assert "VisibleType" in defs + assert "HiddenType" not in defs + + async def test_forward_with_argument_mapping(add_tool): """Test that forward() applies argument mapping correctly."""