Merge pull request #12 from jlowin/context

Add support for request context, progress, logging, etc.

src/fastmcp/__init__.py

@@ -1,6 +1,6 @@
 """FastMCP - A more ergonomic interface for MCP servers."""
 
-from .server import FastMCP
+from .server import FastMCP, Context
 from .utilities.types import Image
 
-__all__ = ["FastMCP", "Image"]
+__all__ = ["FastMCP", "Context", "Image"]

src/fastmcp/resources/types.py

@@ -1,5 +1,6 @@
 """Concrete resource implementations."""
 
+import pydantic_core
 import asyncio
 import json
 from pathlib import Path
@@ -58,8 +59,8 @@ class FunctionResource(Resource):
             if isinstance(result, str):
                 return result
             try:
-                return json.dumps(result, default=pydantic.json.pydantic_encoder)
-            except TypeError:
+                return json.dumps(pydantic_core.to_jsonable_python(result))
+            except (TypeError, pydantic_core.PydanticSerializationError):
                 # If JSON serialization fails, try str()
                 return str(result)
         except Exception as e:

src/fastmcp/server.py

@@ -1,13 +1,20 @@
 """FastMCP - A more ergonomic interface for MCP servers."""
 
+import pydantic_core
+from typing import Any, Literal, Optional, Union
+
+from mcp.server import RequestContext
+from pydantic import BaseModel
+from pydantic.networks import AnyUrl
+
+from fastmcp.utilities.logging import get_logger
 import asyncio
 import functools
 import json
-from typing import Any, Callable, Optional, Sequence, Union, Literal
+from typing import Callable, Sequence
 import inspect
 import re
 
-import pydantic.json
 from mcp.server import Server as MCPServer
 from mcp.server.stdio import stdio_server
 from mcp.server.sse import SseServerTransport
@@ -25,7 +32,7 @@ from fastmcp.exceptions import ResourceError
 from fastmcp.resources import Resource, ResourceManager
 from fastmcp.resources.types import FunctionResource
 from fastmcp.tools import ToolManager
-from fastmcp.utilities.logging import get_logger, configure_logging
+from fastmcp.utilities.logging import configure_logging
 from fastmcp.utilities.types import Image
 
 logger = get_logger(__name__)
@@ -112,12 +119,22 @@ class FastMCP:
             for info in tools
         ]
 
+    def get_context(self) -> Optional["Context"]:
+        try:
+            request_context = self._mcp_server.request_context
+            return Context(request_context=request_context, fastmcp=self)
+        except LookupError:
+            return None
+
     async def call_tool(
         self, name: str, arguments: dict
     ) -> Sequence[Union[TextContent, ImageContent]]:
         """Call a tool by name with arguments."""
         try:
-            result = await self._tool_manager.call_tool(name, arguments)
+            context = self.get_context()
+            result = await self._tool_manager.call_tool(
+                name, arguments, context=context
+            )
             return _convert_to_content(result)
         except Exception as e:
             logger.error(f"Error calling tool {name}: {e}")
@@ -172,13 +189,45 @@ class FastMCP:
         name: Optional[str] = None,
         description: Optional[str] = None,
     ) -> None:
-        """Add a tool to the server."""
+        """Add a tool to the server.
+
+        The tool function can optionally request a Context object by adding a parameter
+        with the Context type annotation. See the @tool decorator for examples.
+
+        Args:
+            func: The function to register as a tool
+            name: Optional name for the tool (defaults to function name)
+            description: Optional description of what the tool does
+        """
         self._tool_manager.add_tool(func, name=name, description=description)
 
     def tool(
         self, name: Optional[str] = None, description: Optional[str] = None
     ) -> Callable:
-        """Decorator to register a tool."""
+        """Decorator to register a tool.
+
+        Tools can optionally request a Context object by adding a parameter with the Context type annotation.
+        The context provides access to MCP capabilities like logging, progress reporting, and resource access.
+
+        Args:
+            name: Optional name for the tool (defaults to function name)
+            description: Optional description of what the tool does
+
+        Example:
+            @server.tool()
+            def my_tool(x: int) -> str:
+                return str(x)
+
+            @server.tool()
+            def tool_with_context(x: int, ctx: Context) -> str:
+                ctx.info(f"Processing {x}")
+                return str(x)
+
+            @server.tool()
+            async def async_tool(x: int, context: Context) -> str:
+                await context.report_progress(50, 100)
+                return str(x)
+        """
         # Check if user passed function directly instead of calling decorator
         if callable(name):
             raise TypeError(
@@ -348,7 +397,7 @@ def _convert_to_content(value: Any) -> Sequence[Union[TextContent, ImageContent]
                 result.append(
                     TextContent(
                         type="text",
-                        text=json.dumps(item, default=pydantic.json.pydantic_encoder),
+                        text=json.dumps(pydantic_core.to_jsonable_python(item)),
                     )
                 )
         return result
@@ -365,6 +414,146 @@ def _convert_to_content(value: Any) -> Sequence[Union[TextContent, ImageContent]
     return [
         TextContent(
             type="text",
-            text=json.dumps(value, indent=2, default=pydantic.json.pydantic_encoder),
+            text=json.dumps(pydantic_core.to_jsonable_python(value)),
         )
     ]
+
+
+class Context(BaseModel):
+    """Context object providing access to MCP capabilities.
+
+    This provides a cleaner interface to MCP's RequestContext functionality.
+    It gets injected into tool and resource functions that request it via type hints.
+
+    To use context in a tool function, add a parameter with the Context type annotation:
+
+    ```python
+    @server.tool()
+    def my_tool(x: int, ctx: Context) -> str:
+        # Log messages to the client
+        ctx.info(f"Processing {x}")
+        ctx.debug("Debug info")
+        ctx.warning("Warning message")
+        ctx.error("Error message")
+
+        # Report progress
+        ctx.report_progress(50, 100)
+
+        # Access resources
+        data = ctx.read_resource("resource://data")
+
+        # Get request info
+        request_id = ctx.request_id
+        client_id = ctx.client_id
+
+        return str(x)
+    ```
+
+    The context parameter name can be anything as long as it's annotated with Context.
+    The context is optional - tools that don't need it can omit the parameter.
+    """
+
+    _request_context: RequestContext
+    _fastmcp: FastMCP
+
+    def __init__(
+        self, *, request_context: RequestContext, fastmcp: FastMCP, **kwargs: Any
+    ):
+        super().__init__(**kwargs)
+        self._request_context = request_context
+        self._fastmcp = fastmcp
+
+    @property
+    def fastmcp(self) -> FastMCP:
+        """Access to the FastMCP server."""
+        return self._fastmcp
+
+    @property
+    def request_context(self) -> RequestContext:
+        """Access to the underlying request context."""
+        return self._request_context
+
+    async def report_progress(
+        self, progress: float, total: Optional[float] = None
+    ) -> None:
+        """Report progress for the current operation.
+
+        Args:
+            progress: Current progress value e.g. 24
+            total: Optional total value e.g. 100
+        """
+
+        progress_token = (
+            self.request_context.meta.progressToken
+            if self.request_context.meta
+            else None
+        )
+
+        if not progress_token:
+            return
+
+        await self.request_context.session.send_progress_notification(
+            progress_token=progress_token, progress=progress, total=total
+        )
+
+    async def read_resource(self, uri: Union[str, AnyUrl]) -> Union[str, bytes]:
+        """Read a resource by URI.
+
+        Args:
+            uri: Resource URI to read
+
+        Returns:
+            The resource content as either text or bytes
+        """
+        return await self._fastmcp.read_resource(uri)
+
+    def log(
+        self,
+        level: Literal["debug", "info", "warning", "error"],
+        message: str,
+        *,
+        logger_name: Optional[str] = None,
+    ) -> None:
+        """Send a log message to the client.
+
+        Args:
+            level: Log level (debug, info, warning, error)
+            message: Log message
+            logger_name: Optional logger name
+            **extra: Additional structured data to include
+        """
+        self.request_context.session.send_log_message(
+            level=level, data=message, logger=logger_name
+        )
+
+    @property
+    def client_id(self) -> Optional[str]:
+        """Get the client ID if available."""
+        return self.request_context.meta.clientId if self.request_context.meta else None
+
+    @property
+    def request_id(self) -> str:
+        """Get the unique ID for this request."""
+        return self.request_context.request_id
+
+    @property
+    def session(self):
+        """Access to the underlying session for advanced usage."""
+        return self.request_context.session
+
+    # Convenience methods for common log levels
+    def debug(self, message: str, **extra: Any) -> None:
+        """Send a debug log message."""
+        self.log("debug", message, **extra)
+
+    def info(self, message: str, **extra: Any) -> None:
+        """Send an info log message."""
+        self.log("info", message, **extra)
+
+    def warning(self, message: str, **extra: Any) -> None:
+        """Send a warning log message."""
+        self.log("warning", message, **extra)
+
+    def error(self, message: str, **extra: Any) -> None:
+        """Send an error log message."""
+        self.log("error", message, **extra)

src/fastmcp/tools.py

@@ -1,12 +1,16 @@
 """Tool management for FastMCP."""
 
 import inspect
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, TYPE_CHECKING
 
 from pydantic import BaseModel, Field, TypeAdapter, validate_call
 
 from .exceptions import ToolError
 from .utilities.logging import get_logger
+import fastmcp
+
+if TYPE_CHECKING:
+    from fastmcp.server import Context
 
 logger = get_logger(__name__)
 
@@ -19,6 +23,9 @@ class Tool(BaseModel):
     description: str = Field(description="Description of what the tool does")
     parameters: dict = Field(description="JSON schema for tool parameters")
     is_async: bool = Field(description="Whether the tool is async")
+    context_kwarg: Optional[str] = Field(
+        None, description="Name of the kwarg that should receive context"
+    )
 
     @classmethod
     def from_function(
@@ -26,6 +33,7 @@ class Tool(BaseModel):
         func: Callable,
         name: Optional[str] = None,
         description: Optional[str] = None,
+        context_kwarg: Optional[str] = None,
     ) -> "Tool":
         """Create a Tool from a function."""
         func_name = name or func.__name__
@@ -39,6 +47,14 @@ class Tool(BaseModel):
         # Get schema from TypeAdapter - will fail if function isn't properly typed
         parameters = TypeAdapter(func).json_schema()
 
+        # Find context parameter if it exists
+        if context_kwarg is None:
+            sig = inspect.signature(func)
+            for param_name, param in sig.parameters.items():
+                if param.annotation is fastmcp.Context:
+                    context_kwarg = param_name
+                    break
+
         # ensure the arguments are properly cast
         func = validate_call(func)
 
@@ -48,11 +64,16 @@ class Tool(BaseModel):
             description=func_doc,
             parameters=parameters,
             is_async=is_async,
+            context_kwarg=context_kwarg,
         )
 
-    async def run(self, arguments: dict) -> Any:
+    async def run(self, arguments: dict, context: Optional["Context"] = None) -> Any:
         """Run the tool with arguments."""
         try:
+            # Inject context if needed
+            if self.context_kwarg and context:
+                arguments[self.context_kwarg] = context
+
             # Call function with proper async handling
             if self.is_async:
                 return await self.func(**arguments)
@@ -92,10 +113,12 @@ class ToolManager:
         self._tools[tool.name] = tool
         return tool
 
-    async def call_tool(self, name: str, arguments: dict) -> Any:
+    async def call_tool(
+        self, name: str, arguments: dict, context: Optional["Context"] = None
+    ) -> Any:
         """Call a tool by name with arguments."""
         tool = self.get_tool(name)
         if not tool:
             raise ToolError(f"Unknown tool: {name}")
 
-        return await tool.run(arguments)
+        return await tool.run(arguments, context=context)

tests/resources/test_function_resources.py

@@ -1,3 +1,4 @@
+from pydantic import BaseModel
 import pytest
 from fastmcp.resources import FunctionResource
 
@@ -80,6 +81,20 @@ class TestFunctionResource:
         with pytest.raises(ValueError, match="Error reading resource function://test"):
             await resource.read()
 
+    async def test_basemodel_conversion(self):
+        """Test handling of BaseModel types."""
+
+        class MyModel(BaseModel):
+            name: str
+
+        resource = FunctionResource(
+            uri="function://test",
+            name="test",
+            func=lambda: MyModel(name="test"),
+        )
+        content = await resource.read()
+        assert content == '{"name": "test"}'
+
     async def test_custom_type_conversion(self):
         """Test handling of custom types."""
 

tests/test_server.py

@@ -1,7 +1,7 @@
 from mcp.shared.memory import (
     create_connected_server_and_client_session as client_session,
 )
-from fastmcp import FastMCP
+from fastmcp import FastMCP, Context
 from fastmcp.resources import FileResource, FunctionResource
 from fastmcp.utilities.types import Image
 from mcp.types import TextContent, ImageContent
@@ -9,7 +9,10 @@ import pytest
 from pydantic import BaseModel
 from pathlib import Path
 import base64
-from typing import Union
+from typing import Union, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fastmcp import Context
 
 
 class TestServer:
@@ -368,3 +371,95 @@ class TestServerResourceTemplates:
         assert isinstance(resource, FunctionResource)
         result = await resource.read()
         assert result == "Data for test"
+
+
+class TestContextInjection:
+    """Test context injection in tools."""
+
+    async def test_context_detection(self):
+        """Test that context parameters are properly detected."""
+        mcp = FastMCP()
+
+        def tool_with_context(x: int, ctx: Context) -> str:
+            return f"Request {ctx.request_id}: {x}"
+
+        tool = mcp._tool_manager.add_tool(tool_with_context)
+        assert tool.context_kwarg == "ctx"
+
+    async def test_context_injection(self):
+        """Test that context is properly injected into tool calls."""
+        mcp = FastMCP()
+
+        def tool_with_context(x: int, ctx: Context) -> str:
+            assert ctx.request_id is not None
+            return f"Request {ctx.request_id}: {x}"
+
+        mcp.add_tool(tool_with_context)
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("tool_with_context", {"x": 42})
+            assert len(result.content) == 1
+            assert "Request" in result.content[0].text
+            assert "42" in result.content[0].text
+
+    async def test_async_context(self):
+        """Test that context works in async functions."""
+        mcp = FastMCP()
+
+        async def async_tool(x: int, ctx: Context) -> str:
+            assert ctx.request_id is not None
+            return f"Async request {ctx.request_id}: {x}"
+
+        mcp.add_tool(async_tool)
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("async_tool", {"x": 42})
+            assert len(result.content) == 1
+            assert "Async request" in result.content[0].text
+            assert "42" in result.content[0].text
+
+    async def test_context_logging(self):
+        """Test that context logging methods work."""
+        mcp = FastMCP()
+
+        def logging_tool(msg: str, ctx: Context) -> str:
+            ctx.debug("Debug message")
+            ctx.info("Info message")
+            ctx.warning("Warning message")
+            ctx.error("Error message")
+            return f"Logged messages for {msg}"
+
+        mcp.add_tool(logging_tool)
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("logging_tool", {"msg": "test"})
+            assert len(result.content) == 1
+            assert "Logged messages for test" in result.content[0].text
+
+    async def test_optional_context(self):
+        """Test that context is optional."""
+        mcp = FastMCP()
+
+        def no_context(x: int) -> int:
+            return x * 2
+
+        mcp.add_tool(no_context)
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("no_context", {"x": 21})
+            assert len(result.content) == 1
+            assert result.content[0].text == "42"
+
+    async def test_context_resource_access(self):
+        """Test that context can access resources."""
+        mcp = FastMCP()
+
+        @mcp.resource("test://data")
+        def test_resource() -> str:
+            return "resource data"
+
+        @mcp.tool()
+        async def tool_with_resource(ctx: Context) -> str:
+            data = await ctx.read_resource("test://data")
+            return f"Read resource: {data}"
+
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("tool_with_resource", {})
+            assert len(result.content) == 1
+            assert "Read resource: resource data" in result.content[0].text