Allow setting type with ArgTransform

src/fastmcp/tools/tool_transform.py

@@ -112,6 +112,7 @@ from mcp.types import EmbeddedResource, ImageContent, TextContent, ToolAnnotatio
 
 from fastmcp.tools.tool import ParsedFunction, Tool
 from fastmcp.utilities.logging import get_logger
+from fastmcp.utilities.types import get_cached_typeadapter
 
 if TYPE_CHECKING:
     pass
@@ -193,6 +194,7 @@ class ArgTransform:
         name: New name for the argument. Use None to keep original name, or ... for no change.
         description: New description for the argument. Use None to remove description, or ... for no change.
         default: New default value for the argument. Use ... for no change.
+        type: New type for the argument. Use ... for no change.
         drop: If True, remove this argument from the transformed tool's schema.
 
     Examples:
@@ -205,16 +207,20 @@ class ArgTransform:
         # Add a default value (makes argument optional)
         ArgTransform(default=42)
 
+        # Change the type
+        ArgTransform(type=str)
+
         # Drop the argument entirely
         ArgTransform(drop=True)
 
         # Combine multiple transformations
-        ArgTransform(name="new_name", description="New desc", default=None)
+        ArgTransform(name="new_name", description="New desc", default=None, type=int)
     """
 
     name: str | None | EllipsisType = ...
     description: str | None | EllipsisType = ...
     default: Any | EllipsisType = ...
+    type: Any | EllipsisType = ...
     drop: bool = False
 
 
@@ -259,9 +265,18 @@ class TransformedTool(Tool):
         """
         from fastmcp.tools.tool import _convert_to_content
 
+        # Fill in missing arguments with schema defaults to ensure
+        # ArgTransform defaults take precedence over function defaults
+        filled_arguments = arguments.copy()
+        properties = self.parameters.get("properties", {})
+
+        for param_name, param_schema in properties.items():
+            if param_name not in filled_arguments and "default" in param_schema:
+                filled_arguments[param_name] = param_schema["default"]
+
         token = _current_tool.set(self)
         try:
-            result = await self.fn(**arguments)
+            result = await self.fn(**filled_arguments)
             return _convert_to_content(result, serializer=self.serializer)
         finally:
             _current_tool.reset(token)
@@ -357,51 +372,20 @@ class TransformedTool(Tool):
                         f"Function declares: {', '.join(sorted(fn_params))}"
                     )
 
-                # The function defines the final schema
-                final_schema = parsed_fn.parameters.copy()
-                # Inherit descriptions from transformed parent where possible
-                fn_props = final_schema.get("properties", {})
-                transformed_props = schema.get("properties", {})
-
-                for param_name in fn_props:
-                    if param_name in transformed_props:
-                        parent_desc = transformed_props[param_name].get("description")
-                        if parent_desc and "description" not in fn_props[param_name]:
-                            fn_props[param_name]["description"] = parent_desc
+                # ArgTransform takes precedence over function signature
+                # Start with function schema as base, then override with transformed schema
+                final_schema = cls._merge_schema_with_precedence(
+                    parsed_fn.parameters, schema
+                )
             else:
                 # With **kwargs, function can access all transformed params
-                # Function params override transformed params if they overlap
+                # ArgTransform takes precedence over function signature
                 # No validation needed - kwargs makes everything accessible
 
-                # Function accepts **kwargs, so use transformed schema as base
-                # and let function override specific parameters
-                fn_props = parsed_fn.parameters.get("properties", {})
-                fn_required = set(parsed_fn.parameters.get("required", []))
-
-                final_props = schema.get("properties", {}).copy()
-                final_required = set(schema.get("required", []))
-
-                # Override with function's parameters
-                for param_name, param_schema in fn_props.items():
-                    # Inherit description from transformed parent if function doesn't provide one
-                    if param_name in final_props and "description" not in param_schema:
-                        param_schema = param_schema.copy()
-                        param_schema["description"] = final_props[param_name].get(
-                            "description"
-                        )
-
-                    final_props[param_name] = param_schema
-
-                    if param_name in fn_required:
-                        final_required.add(param_name)
-                    else:
-                        final_required.discard(param_name)
-
-                final_schema = {
-                    "type": "object",
-                    "properties": final_props,
-                    "required": list(final_required),
-                }
+                # Start with function schema as base, then override with transformed schema
+                final_schema = cls._merge_schema_with_precedence(
+                    parsed_fn.parameters, schema
+                )
 
         # Additional validation: check for naming conflicts after transformation
         if transform_args:
@@ -578,11 +562,76 @@ class TransformedTool(Tool):
             if transform.default is not ...:
                 new_schema["default"] = transform.default
                 is_required = False
+            if transform.type is not ...:
+                # Use TypeAdapter to get proper JSON schema for the type
+                type_schema = get_cached_typeadapter(transform.type).json_schema()
+                # Update the schema with the type information from TypeAdapter
+                new_schema.update(type_schema)
 
             return new_name, new_schema, is_required  # type: ignore[return-value]
 
         raise ValueError(f"Invalid transform: {transform}")
 
+    @staticmethod
+    def _merge_schema_with_precedence(
+        base_schema: dict[str, Any], override_schema: dict[str, Any]
+    ) -> dict[str, Any]:
+        """Merge two schemas, with the override schema taking precedence.
+
+        Args:
+            base_schema: Base schema to start with
+            override_schema: Schema that takes precedence for overlapping properties
+
+        Returns:
+            Merged schema with override taking precedence
+        """
+        merged_props = base_schema.get("properties", {}).copy()
+        merged_required = set(base_schema.get("required", []))
+
+        override_props = override_schema.get("properties", {})
+        override_required = set(override_schema.get("required", []))
+
+        # Override properties
+        for param_name, param_schema in override_props.items():
+            if param_name in merged_props:
+                # Merge the schemas, with override taking precedence
+                base_param = merged_props[param_name].copy()
+                base_param.update(param_schema)
+                merged_props[param_name] = base_param
+            else:
+                merged_props[param_name] = param_schema.copy()
+
+        # Handle required parameters - override takes complete precedence
+        # Start with override's required set
+        final_required = override_required.copy()
+
+        # For parameters not in override, inherit base requirement status
+        # but only if they don't have a default in the final merged properties
+        for param_name in merged_required:
+            if param_name not in override_props:
+                # Parameter not mentioned in override, keep base requirement status
+                final_required.add(param_name)
+            elif (
+                param_name in override_props
+                and "default" not in merged_props[param_name]
+            ):
+                # Parameter in override but no default, keep required if it was required in base
+                if param_name not in override_required:
+                    # Override doesn't specify it as required, and it has no default,
+                    # so inherit from base
+                    final_required.add(param_name)
+
+        # Remove any parameters that have defaults (they become optional)
+        for param_name, param_schema in merged_props.items():
+            if "default" in param_schema:
+                final_required.discard(param_name)
+
+        return {
+            "type": "object",
+            "properties": merged_props,
+            "required": list(final_required),
+        }
+
     @staticmethod
     def _function_has_kwargs(fn: Callable[..., Any]) -> bool:
         """Check if function accepts **kwargs.

tests/tools/test_tool_transform.py

@@ -1,9 +1,10 @@
 import re
-from typing import Annotated, Any
+from dataclasses import dataclass
+from typing import Annotated, Any, TypedDict
 
 import pytest
 from dirty_equals import IsList
-from pydantic import Field
+from pydantic import BaseModel, Field
 from rich import print  # type: ignore
 
 from fastmcp import FastMCP
@@ -388,6 +389,219 @@ async def test_chaining_transformations(add_tool):
     assert "Chained:" in result[0].text  # type: ignore
 
 
+class MyModel(BaseModel):
+    x: int
+    y: str
+
+
+@dataclass
+class MyDataclass:
+    x: int
+    y: str
+
+
+class MyTypedDict(TypedDict):
+    x: int
+    y: str
+
+
+@pytest.mark.parametrize(
+    "py_type, json_type",
+    [
+        (int, "integer"),
+        (float, "number"),
+        (str, "string"),
+        (bool, "boolean"),
+        (list, "array"),
+        (list[int], "array"),
+        (dict, "object"),
+        (dict[str, int], "object"),
+        (MyModel, "object"),
+        (MyDataclass, "object"),
+        (MyTypedDict, "object"),
+    ],
+)
+def test_arg_transform_type_handling(add_tool, py_type, json_type):
+    """Test that ArgTransform type attribute gets applied to schema."""
+    new_tool = Tool.from_tool(
+        add_tool, transform_args={"old_x": ArgTransform(type=py_type)}
+    )
+
+    # Check that the type was changed in the schema
+    x_prop = get_property(new_tool, "old_x")
+    assert x_prop["type"] == json_type
+
+
+def test_arg_transform_annotated_types(add_tool):
+    """Test that ArgTransform works with annotated types and complex types."""
+    from typing import Annotated
+
+    from pydantic import Field
+
+    # Test with Annotated types
+    tool = Tool.from_tool(
+        add_tool,
+        transform_args={
+            "old_x": ArgTransform(
+                type=Annotated[int, Field(description="An annotated integer")]
+            )
+        },
+    )
+
+    x_prop = get_property(tool, "old_x")
+    assert x_prop["type"] == "integer"
+    # The ArgTransform description should override the annotation description
+    # (since we didn't set a description in ArgTransform, it should use the original)
+
+    # Test with Annotated string that has constraints
+    tool2 = Tool.from_tool(
+        add_tool,
+        transform_args={
+            "old_x": ArgTransform(
+                type=Annotated[str, Field(min_length=1, max_length=10)]
+            )
+        },
+    )
+
+    x_prop2 = get_property(tool2, "old_x")
+    assert x_prop2["type"] == "string"
+    assert x_prop2["minLength"] == 1
+    assert x_prop2["maxLength"] == 10
+
+
+def test_arg_transform_precedence_over_function_without_kwargs():
+    """Test that ArgTransform attributes take precedence over function signature (no **kwargs)."""
+
+    @Tool.from_function
+    def base(x: int, y: str = "default") -> str:
+        return f"{x}: {y}"
+
+    # Function signature says x: int with no default, y: str = "function_default"
+    # ArgTransform should override these
+    def custom_fn(x: str = "transform_default", y: int = 99) -> str:
+        return f"custom: {x}, {y}"
+
+    tool = Tool.from_tool(
+        base,
+        transform_fn=custom_fn,
+        transform_args={
+            "x": ArgTransform(type=str, default="transform_default"),
+            "y": ArgTransform(type=int, default=99),
+        },
+    )
+
+    # ArgTransform should take precedence
+    x_prop = get_property(tool, "x")
+    y_prop = get_property(tool, "y")
+
+    assert x_prop["type"] == "string"  # ArgTransform type wins
+    assert x_prop["default"] == "transform_default"  # ArgTransform default wins
+    assert y_prop["type"] == "integer"  # ArgTransform type wins
+    assert y_prop["default"] == 99  # ArgTransform default wins
+
+    # Neither parameter should be required due to ArgTransform defaults
+    assert "x" not in tool.parameters["required"]
+    assert "y" not in tool.parameters["required"]
+
+
+async def test_arg_transform_precedence_over_function_with_kwargs():
+    """Test that ArgTransform attributes take precedence over function signature (with **kwargs)."""
+
+    @Tool.from_function
+    def base(x: int, y: str = "base_default") -> str:
+        return f"{x}: {y}"
+
+    # Function signature has different types/defaults than ArgTransform
+    async def custom_fn(x: str = "function_default", **kwargs) -> str:
+        result = await forward(x=x, **kwargs)
+        return f"custom: {result}"
+
+    tool = Tool.from_tool(
+        base,
+        transform_fn=custom_fn,
+        transform_args={
+            "x": ArgTransform(type=int, default=42),  # Different type and default
+            "y": ArgTransform(description="ArgTransform description"),
+        },
+    )
+
+    # ArgTransform should take precedence
+    x_prop = get_property(tool, "x")
+    y_prop = get_property(tool, "y")
+
+    assert x_prop["type"] == "integer"  # ArgTransform type wins over function's str
+    assert x_prop["default"] == 42  # ArgTransform default wins over function's default
+    assert (
+        y_prop["description"] == "ArgTransform description"
+    )  # ArgTransform description
+
+    # x should not be required due to ArgTransform default
+    assert "x" not in tool.parameters["required"]
+
+    # Test it works at runtime
+    result = await tool.run(arguments={"y": "test"})
+    # Should use ArgTransform default of 42
+    assert "42: test" in result[0].text  # type: ignore
+
+
+def test_arg_transform_combined_attributes():
+    """Test that multiple ArgTransform attributes work together."""
+
+    @Tool.from_function
+    def base(param: int) -> str:
+        return str(param)
+
+    tool = Tool.from_tool(
+        base,
+        transform_args={
+            "param": ArgTransform(
+                name="renamed_param",
+                type=str,
+                description="New description",
+                default="default_value",
+            )
+        },
+    )
+
+    # Check all attributes were applied
+    assert "renamed_param" in tool.parameters["properties"]
+    assert "param" not in tool.parameters["properties"]
+
+    prop = get_property(tool, "renamed_param")
+    assert prop["type"] == "string"
+    assert prop["description"] == "New description"
+    assert prop["default"] == "default_value"
+    assert "renamed_param" not in tool.parameters["required"]  # Has default
+
+
+async def test_arg_transform_type_precedence_runtime():
+    """Test that ArgTransform type changes work correctly at runtime."""
+
+    @Tool.from_function
+    def base(x: int, y: int = 10) -> int:
+        return x + y
+
+    # Transform x to string type but keep same logic
+    async def custom_fn(x: str, y: int = 10) -> str:
+        # Convert string back to int for the original function
+        result = await forward_raw(x=int(x), y=y)
+        # Extract the text from the result
+        result_text = result[0].text
+        return f"String input '{x}' converted to result: {result_text}"
+
+    tool = Tool.from_tool(
+        base, transform_fn=custom_fn, transform_args={"x": ArgTransform(type=str)}
+    )
+
+    # Verify schema shows string type
+    assert get_property(tool, "x")["type"] == "string"
+
+    # Test it works with string input
+    result = await tool.run(arguments={"x": "5", "y": 3})
+    assert "String input '5'" in result[0].text  # type: ignore
+    assert "result: 8" in result[0].text  # type: ignore
+
+
 class TestProxy:
     @pytest.fixture
     def mcp_server(self) -> FastMCP: