From a147eb17f477e3b25a8df693221bde3e1146dd85 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 7 Jun 2025 17:04:21 -0400 Subject: [PATCH] Add tool transformation --- src/fastmcp/tools/__init__.py | 3 +- src/fastmcp/tools/tool.py | 152 ++++--- src/fastmcp/tools/tool_transform.py | 602 ++++++++++++++++++++++++++++ tests/tools/test_tool_transform.py | 421 +++++++++++++++++++ 4 files changed, 1128 insertions(+), 50 deletions(-) create mode 100644 src/fastmcp/tools/tool_transform.py create mode 100644 tests/tools/test_tool_transform.py diff --git a/src/fastmcp/tools/__init__.py b/src/fastmcp/tools/__init__.py index e659a7785..8fa723915 100644 --- a/src/fastmcp/tools/__init__.py +++ b/src/fastmcp/tools/__init__.py @@ -1,4 +1,5 @@ from .tool import Tool, FunctionTool from .tool_manager import ToolManager +from .tool_transform import forward, forward_raw -__all__ = ["Tool", "ToolManager", "FunctionTool"] +__all__ = ["Tool", "ToolManager", "FunctionTool", "forward", "forward_raw"] diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index a62a81306..ca868198c 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -4,6 +4,7 @@ import inspect import json from abc import ABC, abstractmethod from collections.abc import Callable +from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any import pydantic_core @@ -24,7 +25,7 @@ from fastmcp.utilities.types import ( ) if TYPE_CHECKING: - pass + from fastmcp.tools.tool_transform import ArgTransform, TransformedTool logger = get_logger(__name__) @@ -94,6 +95,31 @@ class Tool(FastMCPBaseModel, ABC): """Run the tool with arguments.""" raise NotImplementedError("Subclasses must implement run()") + @classmethod + def from_tool( + cls, + tool: Tool, + transform_fn: Callable[..., Any] | None = None, + name: str | None = None, + transform_args: dict[str, str | ArgTransform | None] | None = None, + description: str | None = None, + tags: set[str] | None = None, + annotations: ToolAnnotations | None = None, + serializer: Callable[[Any], str] | None = None, + ) -> TransformedTool: + from fastmcp.tools.tool_transform import TransformedTool + + return TransformedTool.from_tool( + tool=tool, + transform_fn=transform_fn, + name=name, + transform_args=transform_args, + description=description, + tags=tags, + annotations=annotations, + serializer=serializer, + ) + class FunctionTool(Tool): fn: Callable[..., Any] @@ -110,59 +136,17 @@ class FunctionTool(Tool): serializer: Callable[[Any], str] | None = None, ) -> FunctionTool: """Create a Tool from a function.""" - from fastmcp.server.context import Context - # Reject functions with *args or **kwargs - sig = inspect.signature(fn) - for param in sig.parameters.values(): - if param.kind == inspect.Parameter.VAR_POSITIONAL: - raise ValueError("Functions with *args are not supported as tools") - if param.kind == inspect.Parameter.VAR_KEYWORD: - raise ValueError("Functions with **kwargs are not supported as tools") + parsed_fn = ParsedFunction.from_function(fn, exclude_args=exclude_args) - if exclude_args: - for arg_name in exclude_args: - if arg_name not in sig.parameters: - raise ValueError( - f"Parameter '{arg_name}' in exclude_args does not exist in function." - ) - param = sig.parameters[arg_name] - if param.default == inspect.Parameter.empty: - raise ValueError( - f"Parameter '{arg_name}' in exclude_args must have a default value." - ) - - func_name = name or getattr(fn, "__name__", None) or fn.__class__.__name__ - - if func_name == "": + if name is None and parsed_fn.name == "": raise ValueError("You must provide a name for lambda functions") - func_doc = description or fn.__doc__ - - # if the fn is a callable class, we need to get the __call__ method from here out - if not inspect.isroutine(fn): - fn = fn.__call__ - # if the fn is a staticmethod, we need to work with the underlying function - if isinstance(fn, staticmethod): - fn = fn.__func__ - - type_adapter = get_cached_typeadapter(fn) - schema = type_adapter.json_schema() - - prune_params: list[str] = [] - context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) - if context_kwarg: - prune_params.append(context_kwarg) - if exclude_args: - prune_params.extend(exclude_args) - - schema = compress_schema(schema, prune_params=prune_params) - return cls( - fn=fn, - name=func_name, - description=func_doc, - parameters=schema, + fn=parsed_fn.fn, + name=name or parsed_fn.name, + description=description or parsed_fn.description, + parameters=parsed_fn.parameters, tags=tags or set(), annotations=annotations, serializer=serializer, @@ -217,6 +201,76 @@ class FunctionTool(Tool): return _convert_to_content(result, serializer=self.serializer) +@dataclass +class ParsedFunction: + fn: Callable[..., Any] + name: str + description: str | None + parameters: dict[str, Any] + + @classmethod + def from_function( + cls, + fn: Callable[..., Any], + exclude_args: list[str] | None = None, + validate: bool = True, + ) -> ParsedFunction: + from fastmcp.server.context import Context + + if validate: + sig = inspect.signature(fn) + # Reject functions with *args or **kwargs + for param in sig.parameters.values(): + if param.kind == inspect.Parameter.VAR_POSITIONAL: + raise ValueError("Functions with *args are not supported as tools") + if param.kind == inspect.Parameter.VAR_KEYWORD: + raise ValueError( + "Functions with **kwargs are not supported as tools" + ) + + # Reject exclude_args that don't exist in the function or don't have a default value + if exclude_args: + for arg_name in exclude_args: + if arg_name not in sig.parameters: + raise ValueError( + f"Parameter '{arg_name}' in exclude_args does not exist in function." + ) + param = sig.parameters[arg_name] + if param.default == inspect.Parameter.empty: + raise ValueError( + f"Parameter '{arg_name}' in exclude_args must have a default value." + ) + + # collect name and doc before we potentially modify the function + fn_name = getattr(fn, "__name__", None) or fn.__class__.__name__ + fn_doc = fn.__doc__ + + # if the fn is a callable class, we need to get the __call__ method from here out + if not inspect.isroutine(fn): + fn = fn.__call__ + # if the fn is a staticmethod, we need to work with the underlying function + if isinstance(fn, staticmethod): + fn = fn.__func__ + + type_adapter = get_cached_typeadapter(fn) + schema = type_adapter.json_schema() + + prune_params: list[str] = [] + context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context) + if context_kwarg: + prune_params.append(context_kwarg) + if exclude_args: + prune_params.extend(exclude_args) + + schema = compress_schema(schema, prune_params=prune_params) + return cls( + fn=fn, + name=fn_name, + description=fn_doc, + parameters=schema, + ) + + def _convert_to_content( result: Any, serializer: Callable[[Any], str] | None = None, diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py new file mode 100644 index 000000000..ed3a5d6c0 --- /dev/null +++ b/src/fastmcp/tools/tool_transform.py @@ -0,0 +1,602 @@ +"""# Tool Transformation + +Transform existing tools with modified schemas, argument mappings, and custom behavior. +Use this for creating tool variants, adapting tools for different contexts, or adding +custom logic while preserving the original tool's functionality. + +## Quick Reference + +### Basic Argument Renaming +```python +# Transform specific parent arguments (others pass through unchanged) +new_tool = Tool.from_tool( + original_tool, + transform_args={"old_param": "new_param"} # Only transforms this one arg +) +``` + +### Complex Transformations +```python +from fastmcp.tools.tool_transform import ArgTransform + +new_tool = Tool.from_tool( + original_tool, + transform_args={ + "old_name": ArgTransform(name="new_name", description="Updated desc"), + "unwanted": ArgTransform(drop=True), + "simple": "renamed" + } +) +``` + +### Custom Transform Functions +```python +async def my_transform(new_x: int, new_y: int) -> str: + # Use forward() with transformed argument names + result = await forward(new_x=new_x, new_y=new_y) + return f"Custom: {result}" + +new_tool = Tool.from_tool( + original_tool, + transform_fn=my_transform, + transform_args={"x": "new_x", "y": "new_y"} +) +``` + +### Using **kwargs for Flexibility +```python +async def flexible_transform(**kwargs) -> str: + # kwargs contains all transformed arguments + result = await forward(**kwargs) + return f"Got: {kwargs}" + +new_tool = Tool.from_tool( + original_tool, + transform_fn=flexible_transform, + transform_args={"x": "input_x", "y": "input_y"} +) +``` + +## Key Functions + +- `forward(**kwargs)`: Call parent tool with transformed argument names +- `forward_raw(**kwargs)`: Call parent tool with original argument names + +## Important Notes + +- `transform_args` is optional - if empty/None, all parent arguments pass through unchanged +- Only arguments listed in `transform_args` are transformed, others remain as-is +- Functions with `**kwargs` receive both transformed and untransformed arguments + +## ArgTransform Options + +- `name`: Rename the argument +- `description`: Change the description +- `default`: Add/change default value +- `drop=True`: Remove the argument entirely + +## Common Patterns + +```python +# Chain transformations (partial transforms at each step) +tool1 = Tool.from_tool(original, transform_args={"a": "x"}) # Only transforms 'a' +tool2 = Tool.from_tool(tool1, transform_args={"x": "final"}) # Only transforms 'x' + +# Pure passthrough (no transform_args needed) +enhanced = Tool.from_tool( + original, + name="enhanced_version", + description="Better tool", + tags={"v2", "enhanced"} + # No transform_args = all parent args pass through unchanged +) + +# Drop specific arguments +simplified = Tool.from_tool( + complex_tool, + transform_args={"complex_config": None} # Drops only this arg +) +``` +""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable +from contextvars import ContextVar +from dataclasses import dataclass +from types import EllipsisType +from typing import TYPE_CHECKING, Any + +from mcp.types import EmbeddedResource, ImageContent, TextContent, ToolAnnotations + +from fastmcp.tools.tool import ParsedFunction, Tool +from fastmcp.utilities.logging import get_logger + +if TYPE_CHECKING: + pass + +logger = get_logger(__name__) + + +# Context variable to store current transformed tool +_current_tool: ContextVar[TransformedTool | None] = ContextVar( + "_current_tool", default=None +) + + +async def forward(**kwargs) -> Any: + """Forward to parent tool with argument transformation applied. + + This function can only be called from within a transformed tool's custom + function. It applies argument transformation (renaming, validation) before + calling the parent tool. + + For example, if the parent tool has args `x` and `y`, but the transformed + tool has args `a` and `b`, and an `transform_args` was provided that maps `x` to + `a` and `y` to `b`, then `forward(a=1, b=2)` will call the parent tool with + `x=1` and `y=2`. + + Args: + **kwargs: Arguments to forward to the parent tool (using transformed names). + + Returns: + The result from the parent tool execution. + + Raises: + RuntimeError: If called outside a transformed tool context. + TypeError: If provided arguments don't match the transformed schema. + """ + tool = _current_tool.get() + if tool is None: + raise RuntimeError("forward() can only be called within a transformed tool") + + # Use the forwarding function that handles mapping + return await tool.forwarding_fn(**kwargs) + + +async def forward_raw(**kwargs) -> Any: + """Forward directly to parent tool without transformation. + + This function bypasses all argument transformation and validation, calling the parent + tool directly with the provided arguments. Use this when you need to call the parent + with its original parameter names and structure. + + For example, if the parent tool has args `x` and `y`, then `forward_raw(x=1, + y=2)` will call the parent tool with `x=1` and `y=2`. + + Args: + **kwargs: Arguments to pass directly to the parent tool (using original names). + + Returns: + The result from the parent tool execution. + + Raises: + RuntimeError: If called outside a transformed tool context. + """ + tool = _current_tool.get() + if tool is None: + raise RuntimeError("forward_raw() can only be called within a transformed tool") + + return await tool.parent_tool.run(kwargs) + + +@dataclass(kw_only=True) +class ArgTransform: + """Configuration for transforming a parent tool's argument. + + This class allows fine-grained control over how individual arguments are transformed + when creating a new tool from an existing one. You can rename arguments, change their + descriptions, add default values, or drop them entirely. + + Attributes: + 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. + drop: If True, remove this argument from the transformed tool's schema. + + Examples: + # Rename argument 'old_name' to 'new_name' + ArgTransform(name="new_name") + + # Change description only + ArgTransform(description="Updated description") + + # Add a default value (makes argument optional) + ArgTransform(default=42) + + # Drop the argument entirely + ArgTransform(drop=True) + + # Combine multiple transformations + ArgTransform(name="new_name", description="New desc", default=None) + """ + + name: str | None | EllipsisType = ... + description: str | None | EllipsisType = ... + default: Any | EllipsisType = ... + drop: bool = False + + +class TransformedTool(Tool): + """A tool that is transformed from another tool. + + This class represents a tool that has been created by transforming another tool. + It supports argument renaming, schema modification, custom function injection, + and provides context for the forward() and forward_raw() functions. + + The transformation can be purely schema-based (argument renaming, dropping, etc.) + or can include a custom function that uses forward() to call the parent tool + with transformed arguments. + + Attributes: + parent_tool: The original tool that this tool was transformed from. + fn: The function to execute when this tool is called (either the forwarding + function for pure transformations or a custom user function). + forwarding_fn: Internal function that handles argument transformation and + validation when forward() is called from custom functions. + """ + + parent_tool: Tool + fn: Callable[..., Any] + forwarding_fn: Callable[..., Any] # Always present, handles arg transformation + + async def run( + self, arguments: dict[str, Any] + ) -> list[TextContent | ImageContent | EmbeddedResource]: + """Run the tool with context set for forward() functions. + + This method executes the tool's function while setting up the context + that allows forward() and forward_raw() to work correctly within custom + functions. + + Args: + arguments: Dictionary of arguments to pass to the tool's function. + + Returns: + List of content objects (text, image, or embedded resources) representing + the tool's output. + """ + from fastmcp.tools.tool import _convert_to_content + + token = _current_tool.set(self) + try: + result = await self.fn(**arguments) + return _convert_to_content(result, serializer=self.serializer) + finally: + _current_tool.reset(token) + + @classmethod + def from_tool( + cls, + tool: Tool, + transform_fn: Callable[..., Any] | None = None, + name: str | None = None, + transform_args: dict[str, str | ArgTransform | None] | None = None, + description: str | None = None, + tags: set[str] | None = None, + annotations: ToolAnnotations | None = None, + serializer: Callable[[Any], str] | None = None, + ) -> TransformedTool: + """Create a transformed tool from a parent tool. + + Args: + tool: The parent tool to transform. + transform_fn: Optional custom function. Can use forward() and forward_raw() + to call the parent tool. Functions with **kwargs receive transformed + argument names. + name: New name for the tool. Defaults to parent tool's name. + transform_args: Optional transformations for parent tool arguments. + Only specified arguments are transformed, others pass through unchanged: + - str: Simple rename + - ArgTransform: Complex transformation (rename/description/default/drop) + - None: Drop the argument + description: New description. Defaults to parent's description. + tags: New tags. Defaults to parent's tags. + annotations: New annotations. Defaults to parent's annotations. + serializer: New serializer. Defaults to parent's serializer. + + Returns: + TransformedTool with the specified transformations. + + Examples: + # Transform specific arguments only + Tool.from_tool(parent, transform_args={"old": "new"}) # Others unchanged + + # Custom function with partial transforms + async def custom(x: int, y: int) -> str: + result = await forward(x=x, y=y) + return f"Custom: {result}" + + Tool.from_tool(parent, transform_fn=custom, transform_args={"a": "x", "b": "y"}) + + # Using **kwargs (gets all args, transformed and untransformed) + async def flexible(**kwargs) -> str: + result = await forward(**kwargs) + return f"Got: {kwargs}" + + Tool.from_tool(parent, transform_fn=flexible, transform_args={"a": "x"}) + """ + + # Validate transform_args early + if transform_args: + parent_params = set(tool.parameters.get("properties", {}).keys()) + unknown_args = set(transform_args.keys()) - parent_params + if unknown_args: + raise ValueError( + f"Unknown arguments in transform_args: {', '.join(sorted(unknown_args))}. " + f"Parent tool has: {', '.join(sorted(parent_params))}" + ) + + # Always create the forwarding transform + schema, forwarding_fn = cls._create_forwarding_transform(tool, transform_args) + + if transform_fn is None: + # User wants pure transformation - use forwarding_fn as the main function + final_fn = forwarding_fn + final_schema = schema + else: + # User provided custom function - merge schemas + parsed_fn = ParsedFunction.from_function(transform_fn, validate=False) + final_fn = transform_fn + + has_kwargs = cls._function_has_kwargs(transform_fn) + + # Validate function parameters against transformed schema + fn_params = set(parsed_fn.parameters.get("properties", {}).keys()) + transformed_params = set(schema.get("properties", {}).keys()) + + if not has_kwargs: + # Without **kwargs, function must declare all transformed params + # Check if function is missing any parameters required after transformation + missing_params = transformed_params - fn_params + if missing_params: + raise ValueError( + f"Function missing parameters required after transformation: " + f"{', '.join(sorted(missing_params))}. " + 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 + else: + # With **kwargs, function can access all transformed params + # Function params override transformed params if they overlap + # 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), + } + + # Additional validation: check for naming conflicts after transformation + if transform_args: + new_names = [] + for old_name, transform in transform_args.items(): + if isinstance(transform, str): + new_names.append(transform) + elif isinstance(transform, ArgTransform) and not transform.drop: + if transform.name is not ... and transform.name is not None: + new_names.append(transform.name) + else: + new_names.append(old_name) + + # Check for duplicate names after transformation + name_counts = {} + for arg_name in new_names: + name_counts[arg_name] = name_counts.get(arg_name, 0) + 1 + + duplicates = [ + arg_name for arg_name, count in name_counts.items() if count > 1 + ] + if duplicates: + raise ValueError( + f"Multiple arguments would be mapped to the same names: " + f"{', '.join(sorted(duplicates))}" + ) + + final_description = description if description is not None else tool.description + + return cls( + fn=final_fn, + forwarding_fn=forwarding_fn, + parent_tool=tool, + name=name or tool.name, + description=final_description, + parameters=final_schema, + tags=tags or tool.tags, + annotations=annotations or tool.annotations, + serializer=serializer or tool.serializer, + ) + + @classmethod + def _create_forwarding_transform( + cls, + parent_tool: Tool, + transform_args: dict[str, str | ArgTransform | None] | None, + ) -> tuple[dict[str, Any], Callable[..., Any]]: + """Create schema and forwarding function that encapsulates all transformation logic. + + This method builds a new JSON schema for the transformed tool and creates a + forwarding function that validates arguments against the new schema and maps + them back to the parent tool's expected arguments. + + Args: + parent_tool: The original tool to transform. + transform_args: Dictionary defining how to transform each argument. + + Returns: + A tuple containing: + - dict: The new JSON schema for the transformed tool + - Callable: Async function that validates and forwards calls to the parent tool + """ + + # Build transformed schema and mapping + parent_props = parent_tool.parameters.get("properties", {}).copy() + parent_required = set(parent_tool.parameters.get("required", [])) + + new_props = {} + new_required = set() + new_to_old = {} + + for old_name, old_schema in parent_props.items(): + # Check if parameter is in transform_args + if transform_args and old_name in transform_args: + transform = transform_args[old_name] + else: + transform = ... # Default behavior - pass through + + transform_result = cls._apply_single_transform( + old_name, + old_schema, + transform, + old_name in parent_required, + ) + + if transform_result: + new_name, new_schema, is_required = transform_result + new_props[new_name] = new_schema + new_to_old[new_name] = old_name + if is_required: + new_required.add(new_name) + + schema = { + "type": "object", + "properties": new_props, + "required": list(new_required), + } + + # Create forwarding function that closes over everything it needs + async def _forward(**kwargs): + # Validate arguments + valid_args = set(new_props.keys()) + provided_args = set(kwargs.keys()) + unknown_args = provided_args - valid_args + + if unknown_args: + raise TypeError( + f"Got unexpected keyword argument(s): {', '.join(sorted(unknown_args))}" + ) + + # Check required arguments + missing_args = new_required - provided_args + if missing_args: + raise TypeError( + f"Missing required argument(s): {', '.join(sorted(missing_args))}" + ) + + # Map arguments to parent names + parent_args = {} + for new_name, value in kwargs.items(): + old_name = new_to_old.get(new_name, new_name) + parent_args[old_name] = value + + return await parent_tool.run(parent_args) + + return schema, _forward + + @staticmethod + def _apply_single_transform( + old_name: str, + old_schema: dict[str, Any], + transform: str | ArgTransform | None | EllipsisType, + is_required: bool, + ) -> tuple[str, dict[str, Any], bool] | None: + """Apply transformation to a single parameter. + + This method handles the transformation of a single argument according to + the specified transformation rules. + + Args: + old_name: Original name of the parameter. + old_schema: Original JSON schema for the parameter. + transform: Transformation to apply (string for rename, ArgTransform for complex, + None to drop, ... to pass through unchanged). + is_required: Whether the original parameter was required. + + Returns: + Tuple of (new_name, new_schema, new_is_required) if parameter should be kept, + None if parameter should be dropped. + """ + if transform is ...: + # Not in transform_args - pass through + return old_name, old_schema.copy(), is_required + elif transform is None: + # Explicitly set to None in transform_args - drop the parameter + return None + + if isinstance(transform, str): + # Simple rename + return transform, old_schema.copy(), is_required + + if isinstance(transform, ArgTransform): + if transform.drop: + return None + + if transform.name is not ...: + new_name = transform.name or old_name # Handle None case + else: + new_name = old_name + new_schema = old_schema.copy() + + if transform.description is not ...: + new_schema["description"] = transform.description + if transform.default is not ...: + new_schema["default"] = transform.default + is_required = False + + return new_name, new_schema, is_required # type: ignore[return-value] + + raise ValueError(f"Invalid transform: {transform}") + + @staticmethod + def _function_has_kwargs(fn: Callable[..., Any]) -> bool: + """Check if function accepts **kwargs. + + This determines whether a custom function can accept arbitrary keyword arguments, + which affects how schemas are merged during tool transformation. + + Args: + fn: Function to inspect. + + Returns: + True if the function has a **kwargs parameter, False otherwise. + """ + sig = inspect.signature(fn) + return any( + p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values() + ) diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py new file mode 100644 index 000000000..1067d36bc --- /dev/null +++ b/tests/tools/test_tool_transform.py @@ -0,0 +1,421 @@ +import re +from typing import Annotated, Any + +import pytest +from dirty_equals import IsList +from pydantic import Field +from rich import print # type: ignore + +from fastmcp import FastMCP +from fastmcp.client.client import Client +from fastmcp.tools import Tool, forward, forward_raw +from fastmcp.tools.tool import FunctionTool +from fastmcp.tools.tool_transform import ArgTransform, TransformedTool + + +def get_property(tool: Tool, name: str) -> dict[str, Any]: + return tool.parameters["properties"][name] + + +@pytest.fixture +def add_tool() -> FunctionTool: + def add( + old_x: Annotated[int, Field(description="old_x description")], old_y: int = 10 + ) -> int: + print("running!") + return old_x + old_y + + return Tool.from_function(add) + + +def test_tool_from_tool_no_change(add_tool): + new_tool = Tool.from_tool(add_tool) + assert isinstance(new_tool, TransformedTool) + assert new_tool.parameters == add_tool.parameters + assert new_tool.name == add_tool.name + assert new_tool.description == add_tool.description + + +async def test_tool_change_arg_name_with_string(add_tool): + new_tool = Tool.from_tool(add_tool, transform_args={"old_x": "new_x"}) + + assert sorted(new_tool.parameters["properties"]) == ["new_x", "old_y"] + assert get_property(new_tool, "new_x") == get_property(add_tool, "old_x") + assert get_property(new_tool, "old_y") == get_property(add_tool, "old_y") + assert new_tool.parameters["required"] == ["new_x"] + result = await new_tool.run(arguments={"new_x": 1, "old_y": 2}) + assert result[0].text == "3" # type: ignore + + +async def test_renamed_arg_description_is_maintained(add_tool): + new_tool = Tool.from_tool(add_tool, transform_args={"old_x": "new_x"}) + assert get_property(new_tool, "new_x")["description"] == "old_x description" + + +async def test_tool_defaults_are_maintained_on_unmapped_args(add_tool): + new_tool = Tool.from_tool(add_tool, transform_args={"old_x": "new_x"}) + result = await new_tool.run(arguments={"new_x": 1}) + assert result[0].text == "11" # type: ignore + + +async def test_tool_defaults_are_maintained_on_mapped_args(add_tool): + new_tool = Tool.from_tool(add_tool, transform_args={"old_y": "new_y"}) + result = await new_tool.run(arguments={"old_x": 1}) + assert result[0].text == "11" # type: ignore + + +def test_tool_change_arg_name_with_arg_transform(add_tool): + new_tool = Tool.from_tool( + add_tool, transform_args={"old_x": ArgTransform(name="new_x")} + ) + + assert sorted(new_tool.parameters["properties"]) == ["new_x", "old_y"] + assert get_property(new_tool, "new_x") == get_property(add_tool, "old_x") + assert get_property(new_tool, "old_y") == get_property(add_tool, "old_y") + assert new_tool.parameters["required"] == ["new_x"] + + +def test_tool_change_arg_description(add_tool): + new_tool = Tool.from_tool( + add_tool, transform_args={"old_x": ArgTransform(description="new description")} + ) + assert get_property(new_tool, "old_x")["description"] == "new description" + + +async def test_tool_drop_arg_with_none(add_tool): + # drop the arg with a default value + new_tool = Tool.from_tool(add_tool, transform_args={"old_y": None}) + assert sorted(new_tool.parameters["properties"]) == ["old_x"] + result = await new_tool.run(arguments={"old_x": 1}) + assert result[0].text == "11" # type: ignore + + +async def test_tool_drop_arg_with_arg_transform(add_tool): + new_tool = Tool.from_tool( + add_tool, transform_args={"old_y": ArgTransform(drop=True)} + ) + assert sorted(new_tool.parameters["properties"]) == ["old_x"] + result = await new_tool.run(arguments={"old_x": 1}) + assert result[0].text == "11" # type: ignore + + +async def test_dropped_args_error_if_provided(add_tool): + new_tool = Tool.from_tool( + add_tool, transform_args={"old_y": ArgTransform(drop=True)} + ) + with pytest.raises( + TypeError, match="Got unexpected keyword argument\\(s\\): old_y" + ): + await new_tool.run(arguments={"old_x": 1, "old_y": 2}) + + +async def test_forward_with_argument_mapping(add_tool): + """Test that forward() applies argument mapping correctly.""" + + async def custom_fn(new_x: int, new_y: int = 5) -> int: + return await forward(new_x=new_x, new_y=new_y) + + new_tool = Tool.from_tool( + add_tool, + transform_fn=custom_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + + result = await new_tool.run(arguments={"new_x": 2, "new_y": 3}) + assert result[0].text == "5" # type: ignore + + +async def test_forward_with_incorrect_args_raises_error(add_tool): + async def custom_fn(new_x: int, new_y: int = 5) -> int: + # the forward should use the new args, not the old ones + return await forward(old_x=new_x, old_y=new_y) + + new_tool = Tool.from_tool( + add_tool, + transform_fn=custom_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + with pytest.raises( + TypeError, match=re.escape("Got unexpected keyword argument(s): old_x, old_y") + ): + await new_tool.run(arguments={"new_x": 2, "new_y": 3}) + + +async def test_forward_raw_without_argument_mapping(add_tool): + """Test that forward_raw() calls parent directly without mapping.""" + + async def custom_fn(new_x: int, new_y: int = 5) -> int: + # Call parent directly with original argument names + result = await forward_raw(old_x=new_x, old_y=new_y) + return result + + new_tool = Tool.from_tool( + add_tool, + transform_fn=custom_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + + result = await new_tool.run(arguments={"new_x": 2, "new_y": 3}) + assert result[0].text == "5" # type: ignore + + +async def test_custom_fn_with_kwargs_and_no_transform_args(add_tool): + async def custom_fn(extra: int, **kwargs) -> int: + sum = await forward(**kwargs) + return int(sum[0].text) + extra # type: ignore[attr-defined] + + new_tool = Tool.from_tool(add_tool, transform_fn=custom_fn) + result = await new_tool.run(arguments={"extra": 1, "old_x": 2, "old_y": 3}) + assert result[0].text == "6" # type: ignore + assert new_tool.parameters["required"] == IsList( + "extra", "old_x", check_order=False + ) + assert list(new_tool.parameters["properties"]) == IsList( + "extra", "old_x", "old_y", check_order=False + ) + + +async def test_fn_with_kwargs_passes_through_original_args(add_tool): + async def custom_fn(new_y: int = 5, **kwargs) -> int: + assert kwargs == {"old_y": 3} + result = await forward(old_x=new_y, **kwargs) + return result + + new_tool = Tool.from_tool(add_tool, transform_fn=custom_fn) + result = await new_tool.run(arguments={"new_y": 2, "old_y": 3}) + assert result[0].text == "5" # type: ignore + + +async def test_fn_with_kwargs_receives_transformed_arg_names(add_tool): + """Test that **kwargs receives arguments with their transformed names from transform_args.""" + + async def custom_fn(new_x: int, **kwargs) -> int: + # kwargs should contain 'old_y': 3 (transformed name), not 'old_y': 3 (original name) + assert kwargs == {"old_y": 3} + result = await forward(new_x=new_x, **kwargs) + return result + + new_tool = Tool.from_tool( + add_tool, transform_fn=custom_fn, transform_args={"old_x": "new_x"} + ) + result = await new_tool.run(arguments={"new_x": 2, "old_y": 3}) + assert result[0].text == "5" # type: ignore + + +async def test_fn_with_kwargs_handles_partial_explicit_args(add_tool): + """Test that function can explicitly handle some transformed args while others pass through kwargs.""" + + async def custom_fn(new_x: int, some_other_param: str = "default", **kwargs) -> int: + # x is explicitly handled, y should come through kwargs with transformed name + assert kwargs == {"old_y": 7} + result = await forward(new_x=new_x, **kwargs) + return result + + new_tool = Tool.from_tool( + add_tool, transform_fn=custom_fn, transform_args={"old_x": "new_x"} + ) + result = await new_tool.run( + arguments={"new_x": 3, "old_y": 7, "some_other_param": "test"} + ) + assert result[0].text == "10" # type: ignore + + +async def test_fn_with_kwargs_mixed_mapped_and_unmapped_args(add_tool): + """Test **kwargs behavior with mix of mapped and unmapped arguments.""" + + async def custom_fn(new_x: int, **kwargs) -> int: + # new_x is explicitly handled, old_y should pass through kwargs with original name (unmapped) + assert kwargs == {"old_y": 5} + result = await forward(new_x=new_x, **kwargs) + return result + + new_tool = Tool.from_tool( + add_tool, transform_fn=custom_fn, transform_args={"old_x": "new_x"} + ) # only map 'a' + result = await new_tool.run(arguments={"new_x": 1, "old_y": 5}) + assert result[0].text == "6" # type: ignore + + +async def test_fn_with_kwargs_dropped_args_not_in_kwargs(add_tool): + """Test that dropped arguments don't appear in **kwargs.""" + + async def custom_fn(new_x: int, **kwargs) -> int: + # 'b' was dropped, so kwargs should be empty + assert kwargs == {} + # Can't use 'old_y' since it was dropped, so just use 'old_x' mapped to 'new_x' + result = await forward(new_x=new_x) + return result + + new_tool = Tool.from_tool( + add_tool, + transform_fn=custom_fn, + transform_args={"old_x": "new_x", "old_y": None}, + ) # drop 'old_y' + result = await new_tool.run(arguments={"new_x": 8}) + # 8 + 10 (default value of b in parent) + assert result[0].text == "18" # type: ignore[attr-defined] + + +async def test_forward_outside_context_raises_error(): + """Test that forward() raises RuntimeError when called outside a transformed tool.""" + with pytest.raises( + RuntimeError, + match=re.escape("forward() can only be called within a transformed tool"), + ): + await forward(new_x=1, old_y=2) + + +async def test_forward_raw_outside_context_raises_error(): + """Test that forward_raw() raises RuntimeError when called outside a transformed tool.""" + with pytest.raises( + RuntimeError, + match=re.escape("forward_raw() can only be called within a transformed tool"), + ): + await forward_raw(new_x=1, old_y=2) + + +def test_transform_args_validation_unknown_arg(add_tool): + """Test that transform_args with unknown arguments raises ValueError.""" + with pytest.raises( + ValueError, match="Unknown arguments in transform_args: unknown_param" + ): + Tool.from_tool(add_tool, transform_args={"unknown_param": "new_name"}) + + +def test_transform_args_creates_duplicate_names(add_tool): + """Test that transform_args creating duplicate parameter names raises ValueError.""" + with pytest.raises( + ValueError, + match="Multiple arguments would be mapped to the same names: same_name", + ): + Tool.from_tool( + add_tool, transform_args={"old_x": "same_name", "old_y": "same_name"} + ) + + +def test_transform_args_creates_duplicate_names_with_arg_transform(add_tool): + """Test that transform_args creating duplicate parameter names raises ValueError.""" + with pytest.raises( + ValueError, + match="Multiple arguments would be mapped to the same names: same_name", + ): + Tool.from_tool( + add_tool, + transform_args={ + "old_x": ArgTransform(name="same_name"), + "old_y": "same_name", + }, + ) + + +def test_function_without_kwargs_missing_params(add_tool): + """Test that function without **kwargs must declare all transformed params.""" + + def invalid_fn(new_x: int, non_existent: str) -> str: + return "test" + + with pytest.raises( + ValueError, + match="Function missing parameters required after transformation: new_y", + ): + Tool.from_tool( + add_tool, + transform_fn=invalid_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + + +def test_function_without_kwargs_can_have_extra_params(add_tool): + """Test that function without **kwargs can declare extra params beyond transformed ones.""" + + def valid_fn(new_x: int, new_y: int, extra_param: str = "default") -> str: + return f"{new_x + new_y}: {extra_param}" + + # This should work fine - function declares all required params plus an extra one + tool = Tool.from_tool( + add_tool, + transform_fn=valid_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + + # The final schema should include all function parameters + assert "new_x" in tool.parameters["properties"] + assert "new_y" in tool.parameters["properties"] + assert "extra_param" in tool.parameters["properties"] + + +def test_function_with_kwargs_can_add_params(add_tool): + """Test that function with **kwargs can add new parameters.""" + + async def valid_fn(extra_param: str, **kwargs) -> str: + result = await forward(**kwargs) + return f"{extra_param}: {result}" + + # This should work fine - kwargs allows access to all transformed params + tool = Tool.from_tool( + add_tool, + transform_fn=valid_fn, + transform_args={"old_x": "new_x", "old_y": "new_y"}, + ) + + # extra_param is added, new_x and new_y are available + assert "extra_param" in tool.parameters["properties"] + assert "new_x" in tool.parameters["properties"] + assert "new_y" in tool.parameters["properties"] + + +async def test_chaining_transformations(add_tool): + """Test that transformed tools can be transformed again.""" + # First transformation + tool1 = Tool.from_tool(add_tool, transform_args={"old_x": "x"}) + + # Second transformation on the already-transformed tool + tool2 = Tool.from_tool(tool1, transform_args={"x": "final_x"}) + + # Should work with the final names + result = await tool2.run(arguments={"final_x": 5, "old_y": 3}) + assert result[0].text == "8" # type: ignore + + # And forward() in a custom function should work + async def custom(final_x: int, old_y: int) -> str: + # forward() goes to tool1, which has 'final_x' and 'old_y' after transformation + result = await forward(final_x=final_x, old_y=old_y) + return f"Chained: {result}" + + tool3 = Tool.from_tool(tool1, transform_fn=custom, transform_args={"x": "final_x"}) + + result = await tool3.run(arguments={"final_x": 5, "old_y": 3}) + assert "Chained:" in result[0].text # type: ignore + + +class TestProxy: + @pytest.fixture + def mcp_server(self) -> FastMCP: + mcp = FastMCP() + + @mcp.tool + def add(old_x: int, old_y: int = 10) -> int: + return old_x + old_y + + return mcp + + @pytest.fixture + def proxy_server(self, mcp_server: FastMCP) -> FastMCP: + from fastmcp.client.transports import FastMCPTransport + + proxy = FastMCP.as_proxy(Client(transport=FastMCPTransport(mcp_server))) + return proxy + + async def test_transform_proxy(self, proxy_server: FastMCP): + # when adding transformed tools to proxy servers. Needs separate investigation. + + add_tool = await proxy_server.get_tool("add") + new_add_tool = Tool.from_tool( + add_tool, name="add_transformed", transform_args={"old_x": "new_x"} + ) + proxy_server.add_tool(new_add_tool) + + async with Client(proxy_server) as client: + # The tool should be registered with its transformed name + result = await client.call_tool("add_transformed", {"new_x": 1, "old_y": 2}) + assert result[0].text == "3" # type: ignore