Merge pull request #13 from jlowin/context

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(

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

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)