import json import logging import pytest from pydantic import BaseModel from fastmcp.exceptions import ToolError from fastmcp.tools import ToolManager from fastmcp.tools.tool import Tool class TestAddTools: def test_basic_function(self): """Test registering and running a basic function.""" def add(a: int, b: int) -> int: """Add two numbers.""" return a + b manager = ToolManager() manager.add_tool_from_fn(add) tool = manager.get_tool("add") assert tool is not None assert tool.name == "add" assert tool.description == "Add two numbers." assert tool.is_async is False assert tool.parameters["properties"]["a"]["type"] == "integer" assert tool.parameters["properties"]["b"]["type"] == "integer" async def test_async_function(self): """Test registering and running an async function.""" async def fetch_data(url: str) -> str: """Fetch data from URL.""" return f"Data from {url}" manager = ToolManager() manager.add_tool_from_fn(fetch_data) tool = manager.get_tool("fetch_data") assert tool is not None assert tool.name == "fetch_data" assert tool.description == "Fetch data from URL." assert tool.is_async is True assert tool.parameters["properties"]["url"]["type"] == "string" def test_pydantic_model_function(self): """Test registering a function that takes a Pydantic model.""" class UserInput(BaseModel): name: str age: int def create_user(user: UserInput, flag: bool) -> dict: """Create a new user.""" return {"id": 1, **user.model_dump()} manager = ToolManager() manager.add_tool_from_fn(create_user) tool = manager.get_tool("create_user") assert tool is not None assert tool.name == "create_user" assert tool.description == "Create a new user." assert tool.is_async is False assert "name" in tool.parameters["$defs"]["UserInput"]["properties"] assert "age" in tool.parameters["$defs"]["UserInput"]["properties"] assert "flag" in tool.parameters["properties"] def test_add_invalid_tool(self): manager = ToolManager() with pytest.raises(AttributeError): manager.add_tool_from_fn(1) # type: ignore def test_add_lambda(self): manager = ToolManager() tool = manager.add_tool_from_fn(lambda x: x, name="my_tool") assert tool.name == "my_tool" def test_add_lambda_with_no_name(self): manager = ToolManager() with pytest.raises( ValueError, match="You must provide a name for lambda functions" ): manager.add_tool_from_fn(lambda x: x) def test_warn_on_duplicate_tools(self, caplog): """Test warning on duplicate tools.""" manager = ToolManager(duplicate_behavior="warn") def test_fn(x: int) -> int: return x manager.add_tool_from_fn(test_fn, name="test_tool") manager.add_tool_from_fn(test_fn, name="test_tool") assert "Tool already exists: test_tool" in caplog.text # Should have the tool assert manager.get_tool("test_tool") is not None def test_disable_warn_on_duplicate_tools(self, caplog): """Test disabling warning on duplicate tools.""" def f(x: int) -> int: return x manager = ToolManager(duplicate_behavior="ignore") manager.add_tool_from_fn(f) with caplog.at_level(logging.WARNING): manager.add_tool_from_fn(f) assert "Tool already exists: f" not in caplog.text def test_error_on_duplicate_tools(self): """Test error on duplicate tools.""" manager = ToolManager(duplicate_behavior="error") def test_fn(x: int) -> int: return x manager.add_tool_from_fn(test_fn, name="test_tool") with pytest.raises(ValueError, match="Tool already exists: test_tool"): manager.add_tool_from_fn(test_fn, name="test_tool") def test_replace_duplicate_tools(self): """Test replacing duplicate tools.""" manager = ToolManager(duplicate_behavior="replace") def original_fn(x: int) -> int: return x def replacement_fn(x: int) -> int: return x * 2 manager.add_tool_from_fn(original_fn, name="test_tool") manager.add_tool_from_fn(replacement_fn, name="test_tool") # Should have replaced with the new function tool = manager.get_tool("test_tool") assert tool is not None assert tool.fn.__name__ == "replacement_fn" def test_ignore_duplicate_tools(self): """Test ignoring duplicate tools.""" manager = ToolManager(duplicate_behavior="ignore") def original_fn(x: int) -> int: return x def replacement_fn(x: int) -> int: return x * 2 manager.add_tool_from_fn(original_fn, name="test_tool") result = manager.add_tool_from_fn(replacement_fn, name="test_tool") # Should keep the original tool = manager.get_tool("test_tool") assert tool is not None assert tool.fn.__name__ == "original_fn" # Result should be the original tool assert result.fn.__name__ == "original_fn" class TestToolTags: """Test functionality related to tool tags.""" def test_add_tool_with_tags(self): """Test adding tags to a tool.""" def example_tool(x: int) -> int: """An example tool with tags.""" return x * 2 manager = ToolManager() tool = manager.add_tool_from_fn(example_tool, tags={"math", "utility"}) assert tool.tags == {"math", "utility"} tool = manager.get_tool("example_tool") assert tool is not None assert tool.tags == {"math", "utility"} def test_add_tool_with_empty_tags(self): """Test adding a tool with empty tags set.""" def example_tool(x: int) -> int: """An example tool with empty tags.""" return x * 2 manager = ToolManager() tool = manager.add_tool_from_fn(example_tool, tags=set()) assert tool.tags == set() def test_add_tool_with_none_tags(self): """Test adding a tool with None tags.""" def example_tool(x: int) -> int: """An example tool with None tags.""" return x * 2 manager = ToolManager() tool = manager.add_tool_from_fn(example_tool, tags=None) assert tool.tags == set() def test_list_tools_with_tags(self): """Test listing tools with specific tags.""" def math_tool(x: int) -> int: """A math tool.""" return x * 2 def string_tool(x: str) -> str: """A string tool.""" return x.upper() def mixed_tool(x: int) -> str: """A tool with multiple tags.""" return str(x) manager = ToolManager() manager.add_tool_from_fn(math_tool, tags={"math"}) manager.add_tool_from_fn(string_tool, tags={"string", "utility"}) manager.add_tool_from_fn(mixed_tool, tags={"math", "utility", "string"}) # Check if we can filter by tags when listing tools math_tools = [tool for tool in manager.list_tools() if "math" in tool.tags] assert len(math_tools) == 2 assert {tool.name for tool in math_tools} == {"math_tool", "mixed_tool"} utility_tools = [ tool for tool in manager.list_tools() if "utility" in tool.tags ] assert len(utility_tools) == 2 assert {tool.name for tool in utility_tools} == {"string_tool", "mixed_tool"} def test_import_tools_preserves_tags(self): """Test that importing tools preserves their tags.""" def tagged_tool(x: int) -> int: """A tool with tags.""" return x source_manager = ToolManager() source_manager.add_tool_from_fn(tagged_tool, tags={"test", "example"}) target_manager = ToolManager() target_manager.import_tools(source_manager, "source/") imported_tool = target_manager.get_tool("source/tagged_tool") assert imported_tool is not None assert imported_tool.tags == {"test", "example"} class TestCallTools: async def test_call_tool(self): def add(a: int, b: int) -> int: """Add two numbers.""" return a + b manager = ToolManager() manager.add_tool_from_fn(add) result = await manager.call_tool("add", {"a": 1, "b": 2}) assert result == 3 async def test_call_async_tool(self): async def double(n: int) -> int: """Double a number.""" return n * 2 manager = ToolManager() manager.add_tool_from_fn(double) result = await manager.call_tool("double", {"n": 5}) assert result == 10 async def test_call_tool_with_default_args(self): def add(a: int, b: int = 1) -> int: """Add two numbers.""" return a + b manager = ToolManager() manager.add_tool_from_fn(add) result = await manager.call_tool("add", {"a": 1}) assert result == 2 async def test_call_tool_with_missing_args(self): def add(a: int, b: int) -> int: """Add two numbers.""" return a + b manager = ToolManager() manager.add_tool_from_fn(add) with pytest.raises(ToolError): await manager.call_tool("add", {"a": 1}) async def test_call_unknown_tool(self): manager = ToolManager() with pytest.raises(ToolError): await manager.call_tool("unknown", {"a": 1}) async def test_call_tool_with_list_int_input(self): def sum_vals(vals: list[int]) -> int: return sum(vals) manager = ToolManager() manager.add_tool_from_fn(sum_vals) # Try both with plain list and with JSON list result = await manager.call_tool("sum_vals", {"vals": "[1, 2, 3]"}) assert result == 6 result = await manager.call_tool("sum_vals", {"vals": [1, 2, 3]}) assert result == 6 async def test_call_tool_with_list_str_or_str_input(self): def concat_strs(vals: list[str] | str) -> str: return vals if isinstance(vals, str) else "".join(vals) manager = ToolManager() manager.add_tool_from_fn(concat_strs) # Try both with plain python object and with JSON list result = await manager.call_tool("concat_strs", {"vals": ["a", "b", "c"]}) assert result == "abc" result = await manager.call_tool("concat_strs", {"vals": '["a", "b", "c"]'}) assert result == "abc" result = await manager.call_tool("concat_strs", {"vals": "a"}) assert result == "a" result = await manager.call_tool("concat_strs", {"vals": '"a"'}) assert result == '"a"' async def test_call_tool_with_complex_model(self): from fastmcp import Context class MyShrimpTank(BaseModel): class Shrimp(BaseModel): name: str shrimp: list[Shrimp] x: None def name_shrimp(tank: MyShrimpTank, ctx: Context) -> list[str]: return [x.name for x in tank.shrimp] manager = ToolManager() manager.add_tool_from_fn(name_shrimp) result = await manager.call_tool( "name_shrimp", {"tank": {"x": None, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}}, ) assert result == ["rex", "gertrude"] result = await manager.call_tool( "name_shrimp", {"tank": '{"x": null, "shrimp": [{"name": "rex"}, {"name": "gertrude"}]}'}, ) assert result == ["rex", "gertrude"] class TestToolSchema: async def test_context_arg_excluded_from_schema(self): from fastmcp import Context def something(a: int, ctx: Context) -> int: return a manager = ToolManager() tool = manager.add_tool_from_fn(something) assert "ctx" not in json.dumps(tool.parameters) assert "Context" not in json.dumps(tool.parameters) assert "ctx" not in tool.fn_metadata.arg_model.model_fields class TestContextHandling: """Test context handling in the tool manager.""" def test_context_parameter_detection(self): """Test that context parameters are properly detected in Tool.from_function().""" from fastmcp import Context def tool_with_context(x: int, ctx: Context) -> str: return str(x) manager = ToolManager() tool = manager.add_tool_from_fn(tool_with_context) assert tool.context_kwarg == "ctx" def tool_without_context(x: int) -> str: return str(x) tool = manager.add_tool_from_fn(tool_without_context) assert tool.context_kwarg is None async def test_context_injection(self): """Test that context is properly injected during tool execution.""" from fastmcp import Context, FastMCP def tool_with_context(x: int, ctx: Context) -> str: assert isinstance(ctx, Context) return str(x) manager = ToolManager() manager.add_tool_from_fn(tool_with_context) mcp = FastMCP() ctx = mcp.get_context() result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx) assert result == "42" async def test_context_injection_async(self): """Test that context is properly injected in async tools.""" from fastmcp import Context, FastMCP async def async_tool(x: int, ctx: Context) -> str: assert isinstance(ctx, Context) return str(x) manager = ToolManager() manager.add_tool_from_fn(async_tool) mcp = FastMCP() ctx = mcp.get_context() result = await manager.call_tool("async_tool", {"x": 42}, context=ctx) assert result == "42" async def test_context_optional(self): """Test that context is optional when calling tools.""" from fastmcp import Context def tool_with_context(x: int, ctx: Context | None = None) -> str: return str(x) manager = ToolManager() manager.add_tool_from_fn(tool_with_context) # Should not raise an error when context is not provided result = await manager.call_tool("tool_with_context", {"x": 42}) assert result == "42" async def test_context_error_handling(self): """Test error handling when context injection fails.""" from fastmcp import Context, FastMCP def tool_with_context(x: int, ctx: Context) -> str: raise ValueError("Test error") manager = ToolManager() manager.add_tool_from_fn(tool_with_context) mcp = FastMCP() ctx = mcp.get_context() with pytest.raises(ToolError, match="Error executing tool tool_with_context"): await manager.call_tool("tool_with_context", {"x": 42}, context=ctx) class TestImportTools: def test_import_tools(self): """Test importing tools from one manager to another with a prefix.""" # Setup source manager with tools source_manager = ToolManager() # Create some test tools def tool1_fn(): return "Tool 1 result" def tool2_fn(): return "Tool 2 result" # Add tools to source manager source_manager.add_tool_from_fn( tool1_fn, name="get_data", description="Get some data" ) source_manager.add_tool_from_fn( tool2_fn, name="process_data", description="Process the data" ) # Create target manager target_manager = ToolManager() # Import tools from source to target prefix = "source/" target_manager.import_tools(source_manager, prefix) # Verify tools were imported with prefixes assert "source/get_data" in target_manager._tools assert "source/process_data" in target_manager._tools # Verify the original tools still exist in source manager assert "get_data" in source_manager._tools assert "process_data" in source_manager._tools # Verify the imported tools have the correct descriptions assert target_manager._tools["source/get_data"].description == "Get some data" assert ( target_manager._tools["source/process_data"].description == "Process the data" ) # Verify the tool functions were properly copied # We can't directly compare functions, so we'll check their __name__ attribute assert target_manager._tools["source/get_data"].fn == tool1_fn assert target_manager._tools["source/process_data"].fn == tool2_fn def test_tool_duplicate_behavior(self): """Test the behavior when importing tools with duplicate names.""" # Setup source and target managers source_manager = ToolManager() target_manager = ToolManager() # Add the same tool name to both managers def source_fn(): return "Source result" def target_fn(): return "Target result" source_manager.add_tool_from_fn(source_fn, name="common_tool") target_manager.add_tool_from_fn( target_fn, name="source/common_tool" ) # Pre-create with the prefixed name # Import tools from source to target target_manager.import_tools(source_manager, "source/") # The original tool in the target manager is replaced by the imported one assert target_manager._tools["source/common_tool"].fn == source_fn def test_import_tools_with_multiple_prefixes(self): """Test importing tools from multiple managers with different prefixes.""" # Setup source managers weather_manager = ToolManager() news_manager = ToolManager() # Add tools to source managers def forecast_fn(): return "Weather forecast" def headlines_fn(): return "News headlines" weather_manager.add_tool_from_fn(forecast_fn, name="forecast") news_manager.add_tool_from_fn(headlines_fn, name="headlines") # Create target manager and import from both sources main_manager = ToolManager() main_manager.import_tools(weather_manager, "weather/") main_manager.import_tools(news_manager, "news/") # Verify tools were imported with correct prefixes assert "weather/forecast" in main_manager._tools assert "news/headlines" in main_manager._tools # Verify the tools are accessible and functioning assert ( main_manager._tools["weather/forecast"].fn.__name__ == forecast_fn.__name__ ) assert ( main_manager._tools["news/headlines"].fn.__name__ == headlines_fn.__name__ ) class TestCustomToolNames: """Test adding tools with custom names that differ from their function names.""" def test_add_tool_with_custom_name(self): """Test adding a tool with a custom name parameter using add_tool_from_fn.""" def original_fn(x: int) -> int: return x * 2 manager = ToolManager() tool = manager.add_tool_from_fn(original_fn, name="custom_name") # The tool is stored under the custom name and its .name is also set to custom_name assert manager.get_tool("custom_name") is not None assert tool.name == "custom_name" assert tool.fn.__name__ == "original_fn" # The tool should not be accessible via its original function name assert manager.get_tool("original_fn") is None def test_add_tool_object_with_custom_key(self): """Test adding a Tool object with a custom key using add_tool().""" def fn(x: int) -> int: return x + 1 # Create a tool with a specific name tool = Tool.from_function(fn, name="my_tool") manager = ToolManager() # Store it under a different name manager.add_tool(tool, key="proxy_tool") # The tool is accessible under the key stored = manager.get_tool("proxy_tool") assert stored is not None # But the tool's .name is unchanged assert stored.name == "my_tool" # The tool is not accessible under its original name assert manager.get_tool("my_tool") is None async def test_call_tool_with_custom_name(self): """Test calling a tool added with a custom name.""" def multiply(a: int, b: int) -> int: """Multiply two numbers.""" return a * b manager = ToolManager() manager.add_tool_from_fn(multiply, name="custom_multiply") # Tool should be callable by its custom name result = await manager.call_tool("custom_multiply", {"a": 5, "b": 3}) assert result == 15 # Original name should not be registered with pytest.raises(ToolError): await manager.call_tool("multiply", {"a": 5, "b": 3}) def test_tool_to_mcp_tool(self): """Test that to_mcp_tool uses the key, not the internal name.""" def some_function(x: int) -> int: return x manager = ToolManager() tool = Tool.from_function(some_function, name="api_function") manager.add_tool(tool) mcp_tools = manager.list_mcp_tools() assert len(mcp_tools) == 1 assert mcp_tools[0].name == "api_function" def test_tool_to_mcp_tool_with_custom_key(self): """Test that to_mcp_tool uses the key, not the internal name.""" def some_function(x: int) -> int: return x manager = ToolManager() tool = Tool.from_function(some_function, name="api_function") manager.add_tool(tool, key="custom-key") # When listing tools for MCP, the key should be used mcp_tools = manager.list_mcp_tools() assert len(mcp_tools) == 1 assert mcp_tools[0].name == "custom-key" def test_import_tools_with_custom_names(self): """Test importing tools with custom names.""" def source_fn(x: int) -> int: return x * 2 # Create a source manager with a tool using custom name source_manager = ToolManager() source_manager.add_tool_from_fn(source_fn, name="custom_source") # Import the tools to a target manager with a prefix target_manager = ToolManager() target_manager.import_tools(source_manager, "prefix/") # The tool should be imported with the prefixed custom name assert target_manager.get_tool("prefix/custom_source") is not None assert target_manager.get_tool("prefix/source_fn") is None def test_replace_tool_keeps_original_name(self): """Test that replacing a tool with "replace" keeps the original name.""" def original_fn(x: int) -> int: return x def replacement_fn(x: int) -> int: return x * 2 # Create a manager with REPLACE behavior manager = ToolManager(duplicate_behavior="replace") # Add the original tool original_tool = manager.add_tool_from_fn(original_fn, name="test_tool") assert original_tool.name == "test_tool" # Replace with a new function but keep the same registered name replacement_tool = manager.add_tool_from_fn(replacement_fn, name="test_tool") # The tool object should have been replaced stored_tool = manager.get_tool("test_tool") assert stored_tool is not None assert stored_tool == replacement_tool # The name should still be the same assert stored_tool.name == "test_tool" # But the function is different assert stored_tool.fn.__name__ == "replacement_fn" def test_mcp_tool_name_for_add_tool(self): """Test MCPTool name for add_tool (key != tool.name).""" def fn(x: int) -> int: return x + 1 tool = Tool.from_function(fn, name="my_tool") manager = ToolManager() manager.add_tool(tool, key="proxy_tool") mcp_tools = manager.list_mcp_tools() assert len(mcp_tools) == 1 assert mcp_tools[0].name == "proxy_tool" def test_mcp_tool_name_for_add_tool_from_fn(self): """Test MCPTool name for add_tool_from_fn (key == tool.name).""" def fn(x: int) -> int: return x + 1 manager = ToolManager() manager.add_tool_from_fn(fn, name="custom_fn") mcp_tools = manager.list_mcp_tools() assert len(mcp_tools) == 1 assert mcp_tools[0].name == "custom_fn"