diff --git a/src/fastmcp/server.py b/src/fastmcp/server.py index 217e9b794..19f3c91db 100644 --- a/src/fastmcp/server.py +++ b/src/fastmcp/server.py @@ -119,12 +119,16 @@ class FastMCP: for info in tools ] - def get_context(self) -> Optional["Context"]: + def get_context(self) -> "Context": + """ + Returns a Context object. Note that the context will only be valid + during a request; outside a request, most methods will error. + """ try: request_context = self._mcp_server.request_context - return Context(request_context=request_context, fastmcp=self) except LookupError: - return None + request_context = None + return Context(request_context=request_context, fastmcp=self) async def call_tool( self, name: str, arguments: dict @@ -457,7 +461,11 @@ class Context(BaseModel): _fastmcp: FastMCP def __init__( - self, *, request_context: RequestContext, fastmcp: FastMCP, **kwargs: Any + self, + *, + request_context: RequestContext = None, + fastmcp: FastMCP = None, + **kwargs: Any, ): super().__init__(**kwargs) self._request_context = request_context @@ -466,11 +474,15 @@ class Context(BaseModel): @property def fastmcp(self) -> FastMCP: """Access to the FastMCP server.""" + if self._fastmcp is None: + raise ValueError("Context is not available outside of a request") return self._fastmcp @property def request_context(self) -> RequestContext: """Access to the underlying request context.""" + if self._request_context is None: + raise ValueError("Context is not available outside of a request") return self._request_context async def report_progress( diff --git a/src/fastmcp/tools.py b/src/fastmcp/tools.py index 7c87c3b91..24a09665d 100644 --- a/src/fastmcp/tools.py +++ b/src/fastmcp/tools.py @@ -71,7 +71,7 @@ class Tool(BaseModel): """Run the tool with arguments.""" try: # Inject context if needed - if self.context_kwarg and context: + if self.context_kwarg: arguments[self.context_kwarg] = context # Call function with proper async handling diff --git a/tests/test_tool_manager.py b/tests/test_tool_manager.py index 8818af938..831c96167 100644 --- a/tests/test_tool_manager.py +++ b/tests/test_tool_manager.py @@ -1,6 +1,7 @@ import logging import pytest from pydantic import BaseModel +from typing import Optional from fastmcp.exceptions import ToolError from fastmcp.tools import ToolManager @@ -153,3 +154,84 @@ class TestCallTools: manager = ToolManager() with pytest.raises(ToolError): await manager.call_tool("unknown", {"a": 1}) + + +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(tool_with_context) + assert tool.context_kwarg == "ctx" + + def tool_without_context(x: int) -> str: + return str(x) + + tool = manager.add_tool(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() + tool = manager.add_tool(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() + tool = manager.add_tool(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: Optional[Context] = None) -> str: + return str(x) + + manager = ToolManager() + tool = manager.add_tool(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() + tool = manager.add_tool(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)